[Question] How to deal with large gradients caused by joint_attach constraints in Euler Integration

[JonathanKuelz]

Description

The SemiImplicitIntegrator enforces "joint attachment" by modeling them as spring-damper systems with parameters joint_attach_ke and joint_attach_kd. Sometimes, it is required to set those relatively high, especially in kinematic chains with 2+ degrees of freedom (see #496 ).

This leads to minimal changes in body positions having a huge impact on joint forces. While I didn't encounter any major issues in forward simulation after tuning these parameters to make the simulation run "more or less smoothly", this actually is problematic for the computation of gradients, which will be dominated by these constraining forces.

Solution (?) proposal

Edited: The previous solution proposal doesn't make sense: Removing the constraint force computation from the backpropagation results in a system where we completely ignore the propagation of forces between joints when computing gradients.

Previous Suggestions

While ignoring the joint attachment forces entirely in the backpropagation might not be formally correct, I found that this significantly stabilizes gradient computations for kinematic chains using the Euler integrator. I would therefore suggest implementing a flag that, when True, results in the joint attachment forces not being considered when computing gradients.

In my use case, detaching body_q and body_qd for before launching eval_body_joints works just fine as a hotfix, but this removes more from the computational graph than necessary (for instance, the computation of q) and (if I am not mistaken) would result in severely wrong gradients if joints are set to anything else that JOINT_MODE_FORCE. However, as there seems to be no option to detach the computation of a specific value within a kernel (see also #590 ), the only alternative I see right now would be re-writing significant parts of the eval_body_joints kernel, potentially splitting it in two.

I don't know if ignoring the computation of attachment forces altogether is the best solution to the described problem, but it is the only one I could come up with to make the Euler integrator usable at all with 2+ joints.

Example

I simulate a 2DOF robot arm movement over 1s and compute a loss as the distance between end-effector and a target position at the final frame.

The screenshots show the gradients of the eef y-coordinate over time. The first shows the current behavior of the SemiImplicit integrator, while the second one shows the one after detaching body_q and body_qd. The second one clearly shows the expected behavior: a single nonzero-gradient for the final step that is propagated back until the first state. On the other hand, the gradients for the current implementation are dominated by noise.

This effect also propagates to all other attributes of the state (body_f, body_qd).

[JonathanKuelz]

I changed this from feature request to question because the solution I proposed is a bad idea. However, the original issue still persists - is there a way to handle this better?