[BUG FIX] Fix invalid gradients in differentiable rigid-body simulation.

[duburcqa]

Description

Follow-up to #2842, fixing the gradient defects surfaced by a full audit of the differentiable rigid-body backward pass, plus the maintainability and convention issues found along the way.

Gradient fixes:

• Checkpoint restore (RigidSolver.load_ckpt) reran the cartesian-space update but not the forward velocity pass, while the replayed substeps skip both; every backward primal in the window then carried stale link velocities, injecting an eps-independent error into every multi-step gradient.
• The equality constraint reverse dropped the velocity-product bias Jdot @ qvel folded into aref (CONNECT and WELD position rows) and the WELD rotation rows' quaternion-derivative bias.
• The constraint reverses accumulated d(jac_qvel)/d(vel) once per chain visit, double-counting shared ancestor dofs of same-root pairs.
• The collision constraint reverse ignored impratio in the regularized friction-cone coefficient of efc_D.
• The prismatic forward-kinematics reverse routed the link position gradient into the joint-frame quaternion term, corrupting quaternion gradients of chains with an offset slide joint (2e-2 relative error).
• dofs.pos.grad (position control, joint stiffness) was never consumed; it now folds back into qpos.grad for all joint types.
• Fixed root links' pose gradients were zeroed by the forward-kinematics reverse even though their pose is persistent input state.
• The adjoint conjugate-gradient solve of the constraint backward stopped on an absolute residual threshold and clamped its step coefficients at an absolute epsilon, capping gradient accuracy around 1e-8 relative whenever the upstream gradient is small; it now converges to the round-off floor.
• The input tape keyed tracked setters by method name only, so two same-step calls on distinct dof subsets (arm + gripper) evicted each other. The key now includes the dof subset, canonicalized for slice (only hashable from Python 3.12) and device-tensor indices, and the reverse skips scalar / array-like inputs; all are valid setter inputs that crashed the backward pass.
• Unsupported configurations now raise at build time instead of silently computing wrong gradients: torsional/rolling friction and attach() with requires_grad=True.
• Collision-row count in the joint-limit reverse follows rows_per_contact instead of a hardcoded 4.

Maintainability: the duplicated impedance-derivative, velocity-contraction, and plane-contact-frame blocks collapse into shared funcs; dead code is removed; grad shuffles on the backward hot path go through zero-copy torch views with the kernels as fallback.

Tests: existing tests are strengthened to fail on each defect above (same-root four-bar connect and nested skew-axis weld fixtures, offset-slide chain, joint stiffness, non-default impratio, split-subset force control, larger equality velocities), and the suite is aligned with the testing conventions. All finite-difference tolerances are re-pinned to floors measured on darwin cpu/Metal and linux cpu/CUDA after the fixes, at a x1.5-5 margin over the worst platform: most fp64 checks tighten by 1 to 5 orders of magnitude (the old floors were these bugs, not physics). Two multistep arms whose position loss was constant now read the quaternion instead of passing vacuously.

Related Issue

Follow-up to #2842.
Addresses #2069: the control-force gradient path it reports broken is fixed and validated per-step by the strengthened multistep tests.
Addresses #2150: the forward-kinematics backward pass it exercises is audited and fixed here; the specific Franka scene remains to be confirmed.

Motivation and Context

The audit found gradient errors ranging from 2e-2 relative down to 3e-7 that finite-difference validation had missed because the scenarios or tolerances did not excite them. Wrong gradients silently break any downstream optimization.

How Has This Been / Can This Be Tested?

• pytest tests/grad/ --backend cpu: 46 passed, 10 skipped (Apple Metal reverse-mode skips) on macOS M4 Max; 56/56 with the Metal skip lifted locally against a Quadrants build containing the issue [BUG FIX] better handling saved terrain #805 fix. Full suite also validated on linux CPU (fp64) and CUDA (fp32) with the tightened tolerances.
• Every strengthened test was verified to FAIL on current main and pass with this PR.
• The checkpoint-replay fix was validated by instrumenting the backward replay: replayed substep states now match the forward pass bit-exactly.

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: 7f289cfeaf

ℹ️ 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".

[duburcqa]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 48e1afb11d

ℹ️ 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".

[duburcqa]

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Breezy!

Reviewed commit: eaa815f634

ℹ️ 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".