Skip to main content

Module adjoint

Module adjoint 

Source
Expand description

Discrete trajectory adjoint for semi-implicit Euler rollouts.

§What this computes

For a rollout x_{t+1} = f(x_t, u_t; π, V) (states x = (q, v), open-loop controls u_t, per-body inertia parameters π, collision-mesh vertices V) and a final-state objective J = g(q_T, v_T), one backward pass yields both parameter sensitivities:

  • dJ/dπ — per body, in the canonical 10-vector packing [m, cx, cy, cz, Ixx, Iyy, Izz, Ixy, Ixz, Iyz] (Task: the inertia parameter adjoint that replaces vcad’s finite-difference ∂J/∂p).
  • dJ/dV — per collision-mesh vertex, ∂J/∂x in the body frame (the contact adjoint that plugs into vcad’s surface_gradient).

§How

Reverse over the trajectory, tangent within the step. The adjoint state λ_t = ∂J/∂x_t is backpropagated through the step Jacobians; those Jacobians are read out exactly with tang::Dual lanes through the scalar-generic step of super::step — no finite differences anywhere.

Semi-implicit Euler has the structure v' = v + dt·a, q' = Φ(q, v') with a = ABA(q, v, u; π, V). Differentiating and collecting terms, with w := Φ_v'ᵀ λ_q' + λ_v' (an nv-covector):

λ_q = Φ_qᵀ λ_q' + dt·(∂a/∂q)ᵀ w     λ_v = w + dt·(∂a/∂v)ᵀ w
dJ/dπ += dt·wᵀ·∂a/∂π               dJ/dV += dt·wᵀ·∂a/∂V

For a model of single-DOF joints only, Φ(q, v') = q + dt·v', so Φ_q = I, Φ_v' = dt·I, and this collapses to the familiar w = dt·λ_q' + λ_v', λ_q = λ_q' + dt·(∂a/∂q)ᵀw — and the driver takes that closed form directly, so a single-DOF model pays nothing for this. With a spherical or free joint Φ is a Lie-group step on quaternion coordinates: Φ_q is not the identity and Φ_v' is nq×nv, not square (see super::step’s module docs for the layout). Both blocks are then read out with dual lanes through Φ alone — nq + nv lanes of a few dozen flops each, negligible beside the ABA lanes.

The two state channels are dual lanes: 2·nv seeded evaluations of the scalar-generic step per timestep.

The parameter and contact channels are not — a lane per parameter would make the backward pass cost O(steps · n_params), i.e. no better than finite differences. Both instead ride one linear solve. Write y := M⁻¹w (read out of a single ABA with zero velocity and zero gravity, where the bias and gravity terms vanish and qdd = M⁻¹·ctrl) and let A^y := J y be the body-frame motion field that joint rates y induce (one O(nb) velocity sweep). Then, since M ∂a/∂π = −∂/∂π ID(q, v, a; π) at fixed (q, v, a) and inverse dynamics is the sum over bodies of ⟨A^y_b, I_b A_b + V_b ×* I_b V_b⟩ (RNEA’s backward pass is the adjoint of its acceleration pass):

dJ/dπ[b][k] −= dt·⟨A^y_b, ΔI·A_b + V_b ×* (ΔI·V_b)⟩,   ΔI := ∂I_b/∂π_k
χ_b          = dt·A^y_b                                (= dt·wᵀ·∂a/∂wrench_b)

ΔI is a constant of the model, priced once outside the time loop; V_b and A_b come from one O(nb) sweep at the nominal (q, v, qdd) recorded by the forward rollout. So all 10·nb inertia directions and all 6·nb wrench directions cost two O(nb) sweeps, not 16·nb dual steps.

The vertex channel then contracts χ_b against the local per-vertex wrench Jacobian (3 tiny dual evaluations of the penalty law per vertex), so vertices never cost a full step either. Cost per timestep is O(nv) dual steps + O(1) ABAs + O(nb) sweeps + O(3·N) local evaluations — independent of the parameter count.

§Contract

  • Joint domain: all of them — revolute/prismatic/fixed plus spherical and free. Multi-DOF joints carry quaternion configuration, so q0 must be laid out per super::step::DofLayout (use its neutral_q), and ∂g/∂q is an nq-vector while ∂g/∂v is an nv-vector.
  • Open-loop control: ctrl(t) must not read the state. A state-feedback law would add ∂u/∂x terms this driver does not model.
  • Objective reads the final state only.
  • Contact: the per-vertex ground-penalty model of super::step (see its module docs for the smoothness contract). The gradient is the exact derivative of that forward model.
  • Determinism: pure f64 arithmetic in a fixed order; two runs are bit-identical.

Structs§

AdjointGradients
Everything one backward pass produces.
AdjointRollout
A differentiable rollout: model, optional contact, initial state, and an open-loop control schedule.
ContactSetup
Ground contact configuration for an adjoint rollout: the plane and the collision skins that feel it.
FinalStateObjective
A final-state objective J = g(q_T, v_T) with its analytic gradient.

Functions§

adjoint_rollout_gradient
One nominal rollout forward, one adjoint pass backward: J, dJ/dπ per body, and (when contact is configured) ∂J/∂x per collision-mesh vertex.
rollout_objective
Run the nominal rollout and return the objective value only (the FD oracle for gates, and a cheap primal probe for callers).

Type Aliases§

ObjectiveGradientFn
The analytic gradient of a final-state objective: (∂g/∂q_T, ∂g/∂v_T).