add smooth contact autodifferentation and freejoin zerograd fix 3/3

[mar-yan24]

Smooth Contact Autodifferentiation 3/3

Overview

Third and final PR in the autodiff chain. Depends on PR google-deepmind#1226 and PR (2/3) merging first.

• PR 1/3: smooth dynamics AD (mark/autodifferentiation -> main)
• PR 2/3: constraint solver implicit diff (mark/autodifferentiation2 -> mark/autodifferentiation)
• PR 3/3 (this): smooth contact AD (mark/autodifferentiation3 -> mark/autodifferentiation2)

This phase, makes contact geometry differentiable. During a taped forward pass, the discrete collision results (dist, pos, frame) get overwritten by smooth analytic distance functions Warp can differentiate, and constraint assembly (efc.J, efc.pos) is recomputed on a differentiable path. Custom adjoint kernels then push efc-level gradients back through the KKT system into contact geometry.

Smooth Collision Recomputation

collision_smooth.py replaces the collision pipeline's contact geometry with differentiable versions. It runs after the discrete pipeline and overwrites contact.{dist, pos, frame} with smooth values from analytic distance functions.

Supported geom type pairs:

• plane-sphere, sphere-sphere, sphere-capsule
• capsule-capsule (2 contacts), plane-capsule (2 contacts)

Unsupported types (box, mesh, convex, ellipsoid, SDF) keep their discrete contact values with zero gradient.

Only runs when the Warp tape is active and qpos.requires_grad is set, so non-AD forward paths have no overhead.

EFC-Level Adjoint Kernels

Two new kernels in adjoint.py connect the solver adjoint (Phase 2) to collision geometry:

• _efc_J_grad_kernel: computes adj_efc_J[i,j] = v[j] * efc_force[i] from the KKT adjoint vector
• _efc_pos_grad_kernel: computes adj_efc_pos from adj_efc_aref via the impedance chain rule d(aref)/d(pos) = -k * imp

Both are launched at the end of solver_implicit_adjoint when collision grad fields are active.

Gradient Field Registry

COLLISION_GRAD_FIELDS in grad.py lists the contact/efc arrays that need requires_grad=True for the collision AD chain. _resolve_field handles dotted field paths like "contact.dist" and "efc.J".

Bug Fixes (carried forward)

• Freejoint zerograd (smooth.py): replaced continue in the kinematics kernel's freejoint branch with an is_free integer flag. Warp's AD replay for continue inside dynamic for-loops emits a goto that skips all adjoint code for that iteration, zeroing gradients.
• Types docstrings: added docstrings for Phase 2 solver retained state fields (solver_h, solver_hfactor, solver_Jaref).

Limitations

• Only primitive geom pairs (plane, sphere, capsule). Box, mesh, convex, and SDF contacts pass through with zero gradient.
• Multi-contact pairs (capsule-capsule, plane-capsule) produce 2 contacts each; only the normal-direction efc row gets a smooth efc_pos gradient.
• Flex contacts are not differentiable.
• Float32 precision: FD validation uses eps=1e-3, atol/rtol=1e-3.

Test Plan

• test_enable_backward_module_flags: verifies collision_smooth has enable_backward=True
• test_euler_step_grad_free: full Euler step AD gradient for freejoint+hinge model
• test_enable_grad_all_smooth_fields: all registered fields toggled correctly via _resolve_field
• test_kinematics_grad("free", ...): freejoint kinematics gradient
• test_forward_without_grad_no_error: non-AD path unaffected
• test_diff_step_produces_nonzero_gradients: end to end