[FEATURE] Improve robustness of differentiable rigid body simulation

[SonSang]

Description

This PR improves the robustness of the differentiable rigid body simulation, and add in-depth unit tests to guarantee the differentiability in various scenarios. This PR is necessary to solve optimization problems that leverage the differentiability of the rigid body simulation.

Related Issue

Resolves Genesis-Embodied-AI/Genesis#

Motivation and Context

Previously, the backward pass of rigid body simulation was brittle, and there was not enough unit tests that verify the backward pass in diverse scenarios. Because of that, several edge cases in the backward pass were not caught and they silently failed downstream applications that leverage the differentiability.

How Has This Been / Can This Be Tested?

To test the differentiability of the rigid body simulation, expanded the unit tests in the existing test_grad.py and split them into the 4 different files as shown below.

tests/test_grad_fd.py
tests/test_grad_optim.py
tests/test_grad_utils.py
tests/test_grad_mpm.py

Each of these tests cover:

• test_grad_fd.py: Contains unit tests for verifying the correctness of the (analytical) gradients from the backward pass by comparing them with the gradients that we estimate from finite difference method. We test under various topology (e.g. cartpole, hopper, etc.) to catch as much edge cases as possible. It tests not only for the single step, but also gradients over multiple steps. Also, it tests under various constraint scenarios, such as joint limit or collision.
• test_grad_optim.py: Contains unit tests that optimize a certain loss value by optimizing input values (e.g. velocity, force). These tests serve as the proxy for the downstream optimization tasks based on differentiability.
• test_grad_utils.py: Check utility functions that are used for the backward pass. For now, there is a single test that verifies if the backward function of the Scene works as expected.
• test_grad_mpm.py: Unlike other tests that check rigid body simulation, this file contains tests for MPM solver.

Screenshots (if appropriate):

Checklist:

• I read the CONTRIBUTING document.
• I followed the Submitting Code Changes section of CONTRIBUTING document.
• I tagged the title correctly (including BUG FIX/FEATURE/MISC/BREAKING)
• I updated the documentation accordingly or no change is needed.
• I tested my changes and added instructions on how to test it for reviewers.

• I have added tests to cover my changes.
• All new and existing tests passed.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 2dcd8e8adf

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: e1f7058f7b

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".