Skip to main content

phyz_diff/
contact_adjoint.rs

1//! Trajectory adjoint through the **convex contact solve** — the same contact
2//! model the forward simulator runs.
3//!
4//! # Why this module exists
5//!
6//! [`crate::rollout`]'s adjoint differentiates a self-contained per-vertex
7//! ground-penalty contact model. The forward simulation path
8//! (`phyz::Simulator::step_with_contacts`) does not use that model: it detects
9//! ground contacts with a margin, assembles the convex problem
10//! (`phyz_contact::assemble`), and solves it with the active-set Newton in
11//! `phyz_contact::convex` under MuJoCo-style solref/solimp stabilization. A
12//! gradient of the penalty model is a gradient of *different physics* than the
13//! engine integrates — the third instance of the solver/gradient-disagreement
14//! bug class in this codebase, and the largest.
15//!
16//! This module closes it. [`convex_adjoint_gradient`] runs **exactly** the
17//! forward step `Simulator::step_with_contacts` runs — same detection, same
18//! assembly, same warm-started solve — caches each step's [`phyz_contact::ContactProblem`]
19//! and [`ContactSolution`], and backpropagates through the implicit-function-
20//! theorem sensitivities of the *actual solve the forward pass performed*
21//! ([`phyz_contact::gradient::FixedPointSensitivity`]), at the regimes the
22//! solver converged to.
23//!
24//! # How the backward pass works
25//!
26//! One contacted step is
27//!
28//! ```text
29//! qdd     = ABA(q, v, u; π)                        (no contact forces)
30//! v_free  = v + dt·qdd
31//! f*      = argmin ½ fᵀ(A(q,π) + R(q,depth))f + fᵀ b(q,v_free)   s.t. cone
32//! v'      = v_free + M⁻¹(q,π) Jᵀ(q) f*
33//! q'      = Φ(q, v')                               (integrate_configuration)
34//! ```
35//!
36//! Differentiating: `dv' = dv_free + d(M⁻¹Jᵀ f*)|_{f fixed} + M⁻¹Jᵀ df*`,
37//! and by the IFT at the converged fixed point with the active set held
38//! fixed, `df*` follows from the exact map linearization
39//! ([`FixedPointSensitivity::apply`]): the differential of the stationarity
40//! residual `(A + R)f* + b − e_n·bias` (impulses held fixed) on the sticking
41//! rows and sliding normal rows, plus the sliding tangential channel
42//! (slip-direction rotation) through each sliding contact's own tangential
43//! block. Separating contacts respond to nothing, by construction.
44//!
45//! The smooth factors (`ABA`, kinematics, assembly, `Φ`) are differentiated by
46//! **central finite differences per lane** — the same discipline
47//! [`crate::semi_implicit_step_jacobians`] documents for its ABA block. The
48//! contact solve itself is *never* finite-differenced: its derivative is the
49//! exact IFT sensitivity of the converged solution, which is what makes the
50//! gradient the derivative of the physics rather than of the solver
51//! iterations. Within a lane the contact set is **frozen** to the identities
52//! the nominal step detected (each contact is re-evaluated from its material
53//! anchor on the body), so a lane never re-runs detection and cannot flip the
54//! active set.
55//!
56//! # Where the gradient is exact, one-sided, or undefined
57//!
58//! - **Interior of a regime** (a contact firmly separated beyond the margin,
59//!   firmly loaded and sticking, or firmly sliding): the gradient is the exact
60//!   derivative of the regularized contact model, up to central-difference
61//!   error (`~1e-9` relative) in the smooth blocks.
62//! - **Contact activation**: the margin band (see
63//!   `ContactMaterial::impedance_at`) ramps the impedance smoothly to zero at
64//!   the band's outer edge, so a contact entering or leaving the *detected*
65//!   set does so at zero force and the gradient is continuous across the
66//!   crossing. This is precisely what the margin was added for. The exception
67//!   is a manifold truncated at `MAX_MANIFOLD_POINTS`: a point excluded by the
68//!   cap can still carry impedance, and the gradient is one-sided there.
69//! - **Stick/slide and load/unload transitions**: the true derivative does not
70//!   exist at the exact transition; the IFT holds the regime the solver
71//!   converged to, so the reported gradient is the one-sided derivative of
72//!   that branch. This matches `phyz_contact::gradient`'s documented contract.
73//! - **Impact timing**: at a genuine contact-making event the gradient of the
74//!   regularized model is biased relative to ideal rigid contact (Suh et al.,
75//!   ICML 2022); see `phyz_contact::gradient`'s module docs.
76//!
77//! # Failure policy
78//!
79//! A step whose contact solve did not converge is not a fixed point and has
80//! no IFT sensitivity; [`FixedPointSensitivity::at`] refuses it by design.
81//! This adjoint **degrades loudly**: it returns
82//! [`ConvexAdjointError::Unconverged`] (or
83//! [`ConvexAdjointError::SingularKkt`] for a degenerate active set) rather
84//! than substituting anything. A silent wrong gradient is the failure mode
85//! this module exists to eliminate.
86//!
87//! # Scope
88//!
89//! Ground **and body-body** contacts. A body-body contact is frozen as a
90//! feature *pair* (`Anchor::Pair`): both surface points in their own body's
91//! frame, the normal in the frame of the body owning the reference face, and
92//! the contact point riding the *other* body's vertex — so a lane sees the
93//! contact rotate and translate with both bodies. That is what carries the
94//! tangential-friction-between-two-moving-bodies channel — a foot on a board's
95//! grip tape, a hand on a lever — which is the case a ground-only adjoint
96//! cannot express at all.
97//!
98//! The per-vertex `∂J/∂(mesh vertex)` channel of [`crate::rollout`] is not
99//! reproduced here; that adjoint remains the (clearly documented) penalty-model
100//! path for the vcad surface-gradient seam.
101
102use phyz_collision::Collision;
103use phyz_contact::gradient::{FixedPointSensitivity, friction_sensitivity};
104use phyz_contact::{
105    ContactAssembly, ContactCache, ContactMaterial, ContactSolution, ContactSolverConfig,
106    GroundSupport, assemble, cylinder_rim_dir, find_contacts,
107    find_ground_contacts_model_with_support, regularization_diag, solve_contacts_warm,
108};
109use phyz_math::{DVec, Vec3};
110use phyz_model::{Model, State};
111use phyz_rigid::{
112    aba, forward_kinematics, integrate_configuration, rotate_free_joint_velocities,
113    strip_free_joint_coriolis,
114};
115
116use crate::rollout::FinalStateObjective;
117pub use crate::rollout::N_INERTIA_PARAMS;
118use crate::rollout::inertia_params;
119
120/// Central-difference step for the smooth (non-solve) blocks.
121///
122/// Smaller than the usual `1e-6` because the margin band's impedance
123/// smoothstep has curvature on the `margin ≈ 1e-3 m` length scale: at
124/// `h = 1e-6` the truncation error of a lane through the band is a few parts
125/// in `1e5` of the step Jacobian, while at `1e-8` it is below `1e-8`. The
126/// pieces being differenced are smooth closed-form evaluations (no iterative
127/// solve inside), so the round-off floor at this step size is `~1e-8`
128/// relative — the truncation/round-off crossover sits near here.
129const FD_EPS: f64 = 1e-8;
130
131/// A differentiable rollout through the forward simulator's contact path.
132///
133/// The forward dynamics are exactly `phyz::Simulator::step_with_contacts`:
134/// ground contact detection with margin, convex assembly, warm-started
135/// active-set Newton solve, semi-implicit Euler integration.
136pub struct ConvexContactRollout<'a> {
137    /// Topology, gravity, timestep, damping, geometry and **nominal** inertias.
138    pub model: &'a Model,
139    /// World z of the ground plane (m).
140    pub ground_height: f64,
141    /// Contact material for every body (the same single-material convention
142    /// `Simulator::step_with_contacts` uses).
143    pub material: ContactMaterial,
144    /// Solver configuration. Use [`ContactSolverConfig::simulation`] to
145    /// differentiate the physics the simulator ships, or
146    /// [`ContactSolverConfig::gradients`] for smoother sensitivities.
147    pub config: ContactSolverConfig,
148    /// Initial joint positions (`Model` layout, length `nq`).
149    pub q0: DVec,
150    /// Initial joint velocities (length `nv`).
151    pub v0: DVec,
152    /// Number of steps.
153    pub steps: usize,
154    /// Open-loop control at step `t` (length `nv`; must not read the state).
155    pub ctrl: &'a dyn Fn(usize) -> DVec,
156}
157
158/// Everything one unified backward pass produces.
159pub struct ConvexAdjointGradients {
160    /// The objective at the nominal rollout.
161    pub objective: f64,
162    /// `dJ/dq0` (length `nq`).
163    pub d_q0: DVec,
164    /// `dJ/dv0` (length `nv`).
165    pub d_v0: DVec,
166    /// `dJ/du_t` per step (each length `nv`).
167    pub d_ctrl: Vec<DVec>,
168    /// `dJ/dπ` per body, canonical packing `[m, cx, cy, cz, Ixx, Iyy, Izz,
169    /// Ixy, Ixz, Iyz]`.
170    pub d_inertia: Vec<[f64; N_INERTIA_PARAMS]>,
171    /// `dJ/dμ` — the friction coefficient of [`ConvexContactRollout::material`].
172    ///
173    /// One scalar, not one per contact, because the rollout applies a single
174    /// material to every body; this is the total derivative with respect to
175    /// that shared coefficient.
176    ///
177    /// **Only sliding contacts contribute.** A sticking contact sits strictly
178    /// inside the cone, so moving the cone boundary does not move the solution,
179    /// and a separating one carries no impulse. So a trajectory that never
180    /// slips reports exactly `0` here — a real property of Coulomb friction
181    /// rather than a missing term, and one a finite difference reproduces
182    /// exactly.
183    pub d_friction: f64,
184    /// `dJ/de` — the restitution coefficient of
185    /// [`ConvexContactRollout::material`], likewise a single shared scalar.
186    ///
187    /// Zero for any trajectory whose contacts all approach below
188    /// [`ContactSolverConfig::restitution_threshold`], where the low-speed ramp
189    /// has taken the effective restitution to zero — again a genuine flat
190    /// region of the model, not an omission.
191    pub d_restitution: f64,
192}
193
194/// Why an adjoint pass refused to produce a gradient.
195///
196/// Every variant is a *loud* refusal in place of a silently wrong number.
197#[derive(Debug, Clone, PartialEq)]
198pub enum ConvexAdjointError {
199    /// A step's contact solve did not converge, so its solution is not a KKT
200    /// point and the IFT does not apply.
201    Unconverged {
202        /// The offending step index.
203        step: usize,
204        /// Iterations the solve used before giving up.
205        iterations: usize,
206        /// The residual it stalled at.
207        residual: f64,
208    },
209    /// The KKT matrix at a step's converged active set is singular — the
210    /// derivative genuinely does not exist there.
211    SingularKkt {
212        /// The offending step index.
213        step: usize,
214    },
215    /// Formerly: the trajectory produced a body-body contact, which this
216    /// adjoint did not cover.
217    ///
218    /// Body-body contacts are differentiated now — see `Anchor::Pair` — and
219    /// nothing constructs this variant any more. It is retained for one release
220    /// so that a downstream `match` on [`ConvexAdjointError`] does not break,
221    /// and because a caller that special-cased this refusal (falling back to a
222    /// sampling optimizer, say) should get a compile-time nudge to delete that
223    /// branch rather than silently keeping dead code that will never fire.
224    #[deprecated(note = "body-body contacts are supported; this variant is never \
225                constructed and will be removed in the next minor release")]
226    BodyBodyContact {
227        /// The step at which the body-body contact appeared.
228        step: usize,
229    },
230    /// A detected contact could not be pinned to a material anchor on its
231    /// body (unsupported geometry variant).
232    UnsupportedGeometry {
233        /// The offending step index.
234        step: usize,
235        /// The body carrying the unsupported geometry.
236        body: usize,
237    },
238}
239
240impl core::fmt::Display for ConvexAdjointError {
241    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
242        match self {
243            Self::Unconverged {
244                step,
245                iterations,
246                residual,
247            } => write!(
248                f,
249                "contact solve at step {step} did not converge \
250                 ({iterations} iterations, residual {residual:.3e}); \
251                 refusing to differentiate a non-KKT point"
252            ),
253            Self::SingularKkt { step } => write!(
254                f,
255                "KKT matrix at step {step} is singular at the converged \
256                 active set; the derivative does not exist there"
257            ),
258            #[allow(deprecated)]
259            Self::BodyBodyContact { step } => write!(
260                f,
261                "body-body contact at step {step}: this refusal is retired — \
262                 body-body contacts are differentiated"
263            ),
264            Self::UnsupportedGeometry { step, body } => write!(
265                f,
266                "could not anchor a ground contact on body {body} at step \
267                 {step}: unsupported geometry for the frozen-contact lanes"
268            ),
269        }
270    }
271}
272
273impl std::error::Error for ConvexAdjointError {}
274
275// ---------------------------------------------------------------------------
276// Frozen contact identities
277// ---------------------------------------------------------------------------
278
279/// A contact pinned to the feature pair that produced it, so a perturbed
280/// configuration re-evaluates *the same* contact smoothly instead of re-running
281/// detection.
282///
283/// Freezing the feature pair is what makes the FD lanes legitimate. Narrow-phase
284/// feature selection (which face is the reference face, which clipped vertices
285/// survive, which box corner is deepest) is combinatorial: a lane that re-ran
286/// detection could return a manifold with a different *number* of points, and
287/// differencing two different contact sets is not a derivative of anything. The
288/// design doc sanctions exactly this — §4.4, "feature selection is treated as a
289/// discrete decision held fixed for the step" — and it is what MuJoCo and Dojo
290/// do. The price is stated there and again in this module's header: the gradient
291/// does not see "the manifold would have had a different point if the box had
292/// rotated slightly more".
293#[derive(Debug, Clone, Copy)]
294enum Anchor {
295    /// A contact against the world `z = ground_height` plane.
296    ///
297    /// The support point decomposes as
298    /// `p_world(q) = pos(q) + Rᵀ(q)·material_point + world_offset`:
299    /// `material_point` is a body-frame point (a box corner, a mesh vertex, a
300    /// capsule hemisphere centre, a sphere centre), and `world_offset` is the
301    /// part that does *not* ride the body: `−r·ẑ` for a sphere or capsule
302    /// dropping onto a z-plane, zero for a box corner or mesh vertex.
303    /// Detection reports which it is per contact
304    /// (`find_ground_contacts_model_with_support`).
305    Ground {
306        body: usize,
307        material_point: Vec3,
308        world_offset: Vec3,
309    },
310    /// A contact against the ground plane on a cylinder's rim.
311    ///
312    /// `p_world(q) = pos(q) + Rᵀ(q)·cap_centre + radius · dir(â(q))`, where
313    /// `â(q) = Rᵀ(q)·axis_local` is the cylinder's axis and `dir` is the rim
314    /// direction `quarter_turns` × 90° around from the barrel's steepest
315    /// downhill `û` (`phyz_contact::cylinder_rim_dir`).
316    ///
317    /// This is [`Anchor::Ground`] with the offset recomputed rather than
318    /// frozen, and the difference is a whole gradient channel. Spinning the
319    /// wheel about its axle leaves `â` alone, so the two agree exactly there —
320    /// the rolling case. *Banking* the wheel rotates `â`, which slides the
321    /// contact around the rim by `r·dû`, and a frozen offset reports none of
322    /// it: measured on a wheel rolling at 0.25 rad of tilt, the frozen form
323    /// gave `dJ/d(tilt) = +7.8e-3` against a finite difference of `−9.5e-3`.
324    /// Wrong sign, not merely small — see
325    /// `phyz-diff/tests/cylinder_contact_adjoint.rs`.
326    ///
327    /// A cylinder standing on its cap takes [`Anchor::Ground`] instead: there
328    /// `û` does not exist, the rim point taken along the shape's own x/y *is*
329    /// a material point, and `Material` is the exact description rather than
330    /// an approximation of one.
331    GroundCylinderRim {
332        body: usize,
333        /// Cap centre, in body coordinates.
334        cap_centre: Vec3,
335        /// The cylinder's axis, in body coordinates.
336        axis_local: Vec3,
337        radius: f64,
338        quarter_turns: u8,
339        /// The world offset detection resolved, used only if a perturbed axis
340        /// lands inside the vertical degeneracy where `û` does not exist.
341        frozen_offset: Vec3,
342    },
343    /// A contact between two moving bodies.
344    ///
345    /// Both surface points are frozen in their own body's frame, and the normal
346    /// is frozen in the frame of whichever body *owns the reference face*
347    /// (`normal_frame`), so that rotating that body rotates the contact normal
348    /// with it. That `dn/dq` channel is the one carrying "the board tilted, so
349    /// the friction direction rotated" — the whole reason a foot-on-grip-tape
350    /// gradient is worth having, and the channel a world-frozen normal drops
351    /// entirely.
352    ///
353    /// Freezing the normal in *a* body frame rather than in the world is what
354    /// generalizes [`Anchor::Ground`] rather than merely extending it: for a
355    /// ground contact the owner is the world, whose frame does not move, so the
356    /// world-fixed `+ẑ` of the ground case is this same rule evaluated on a
357    /// static owner.
358    Pair {
359        body_i: usize,
360        body_j: usize,
361        /// Surface point on `body_i`, in `body_i`'s frame.
362        point_i: Vec3,
363        /// Surface point on `body_j`, in `body_j`'s frame.
364        point_j: Vec3,
365        /// The body whose frame `normal_local` is expressed in — see
366        /// [`Anchor::reference_body`].
367        normal_frame: usize,
368        /// Contact normal in `normal_frame`'s body frame. In world terms it is
369        /// the direction `body_i` must move to separate, matching
370        /// `Collision::contact_normal`.
371        normal_local: Vec3,
372        /// For a curved body (a sphere) the surface point is not a material
373        /// point: it is the centre displaced by the radius along the *live*
374        /// normal, exactly as [`Anchor::Ground`]'s `world_offset` treats a
375        /// ball on the plane. Freezing a sphere's surface point in its own
376        /// frame instead makes the frozen point roll around with the ball, and
377        /// a resting marble on a fixed plate reported `dJ/dx₀ = 102` where
378        /// translation invariance says `1`. `Some(r)` marks that body as a
379        /// sphere of radius `r` whose `point_*` is its centre.
380        radius_i: Option<f64>,
381        radius_j: Option<f64>,
382    },
383}
384
385/// The body's collision shape as a single sphere, if that is what it is.
386fn sphere_radius(model: &Model, body: usize) -> Option<(f64, Vec3)> {
387    use phyz_model::Geometry;
388    let b = model.bodies.get(body)?;
389    if b.collisions.is_empty() {
390        match b.geometry {
391            Some(Geometry::Sphere { radius }) => Some((radius, Vec3::zeros())),
392            _ => None,
393        }
394    } else if b.collisions.len() == 1 {
395        match b.collisions[0].geometry {
396            Geometry::Sphere { radius } => Some((radius, b.collisions[0].origin.pos)),
397            _ => None,
398        }
399    } else {
400        None
401    }
402}
403
404impl Anchor {
405    /// Recover the anchor of a detected collision.
406    ///
407    /// `support_kind` is how the ground contact attaches to its body, as
408    /// reported by `find_ground_contacts_model_with_support`. Detection reports
409    /// it per contact because a multi-shape body no longer determines the
410    /// producing shape by itself. It is unread for a body-body contact, whose
411    /// two surface points are recovered from the manifold geometry instead.
412    ///
413    /// Detection reports the *midsurface* point in both cases. For the ground
414    /// the support point itself is at `z = ground_height − depth`; for a pair
415    /// the two surface points straddle the midsurface by `depth/2` along the
416    /// normal, `body_i`'s on the far side since `+normal` is the direction `i`
417    /// must move to separate.
418    fn of(
419        c: &Collision,
420        support_kind: GroundSupport,
421        state: &State,
422        ground_height: f64,
423        model: &Model,
424    ) -> Self {
425        // `xform.rot` is world→body, so body coordinates of a world point are
426        // `R (p − pos)`, and `Rᵀ` carries a body direction back to world.
427        if c.is_world_j() {
428            let xform = &state.body_xform[c.body_i];
429            let support = Vec3::new(
430                c.contact_point.x,
431                c.contact_point.y,
432                ground_height - c.penetration_depth,
433            );
434            let world_offset = support_kind.world_offset();
435            let material_point = xform.rot * (support - xform.pos - world_offset);
436            if let GroundSupport::CylinderRim {
437                axis,
438                radius,
439                quarter_turns,
440                offset,
441            } = support_kind
442            {
443                return Self::GroundCylinderRim {
444                    body: c.body_i,
445                    cap_centre: material_point,
446                    axis_local: xform.rot * axis,
447                    radius,
448                    quarter_turns,
449                    frozen_offset: offset,
450                };
451            }
452            return Self::Ground {
453                body: c.body_i,
454                material_point,
455                world_offset,
456            };
457        }
458
459        let n = c.contact_normal;
460        let half = 0.5 * c.penetration_depth;
461        // `i` must move along `+n` to separate, so its surface point is the one
462        // on the `−n` side of the midsurface.
463        let surface_i = c.contact_point - n * half;
464        let surface_j = c.contact_point + n * half;
465        let xi = &state.body_xform[c.body_i];
466        let xj = &state.body_xform[c.body_j];
467        let normal_frame = Self::reference_body(c.body_i, c.body_j, n, state, model);
468        let owner = &state.body_xform[normal_frame];
469        // A sphere's anchor is its centre; the surface point is re-derived from
470        // the live normal at evaluation time.
471        let sph_i = sphere_radius(model, c.body_i);
472        let sph_j = sphere_radius(model, c.body_j);
473        Self::Pair {
474            body_i: c.body_i,
475            body_j: c.body_j,
476            point_i: sph_i
477                .map(|(_, centre)| centre)
478                .unwrap_or(xi.rot * (surface_i - xi.pos)),
479            point_j: sph_j
480                .map(|(_, centre)| centre)
481                .unwrap_or(xj.rot * (surface_j - xj.pos)),
482            normal_frame,
483            normal_local: owner.rot * n,
484            radius_i: sph_i.map(|(r, _)| r),
485            radius_j: sph_j.map(|(r, _)| r),
486        }
487    }
488
489    /// Which of the two bodies owns the face the contact normal came from.
490    ///
491    /// The narrow phase does not report this: `Manifold::normal` arrives from
492    /// GJK/EPA as a world direction, a function of *both* poses, with no record
493    /// of which shape's feature generated it. But for the face contacts that
494    /// dominate a resting or riding manifold, the normal *is* a face normal of
495    /// exactly one of the two boxes — which means it is exactly a coordinate
496    /// axis of that body's frame, and merely some oblique direction in the
497    /// other's. So the owner is recoverable after the fact: express the normal
498    /// in each body frame and take the one that lands closest to a principal
499    /// axis.
500    ///
501    /// This matters, and picking the wrong body is not a small error. Freezing
502    /// the normal to the rider instead of to the plank it stands on makes
503    /// `dJ/d(plank roll)` come out as zero when it is genuinely nonzero — the
504    /// tilt no longer steers the friction — which is a *silently* missing
505    /// gradient channel, the failure mode this module exists to prevent. It was
506    /// caught by `body_body_adjoint.rs`'s slide case (adjoint `5e-12` against
507    /// an FD of `-9.7e-3`) and is why that test sweeps every DOF of both bodies
508    /// rather than only the ones the scene obviously moves.
509    ///
510    /// For curved pairs (sphere/capsule) the normal follows from the centres
511    /// and belongs to neither face; the alignment score is then near-arbitrary
512    /// and so is the choice. That is an approximation of the same order as
513    /// freezing the feature pair at all (§4.4), not an additional one: either
514    /// body's frame transports a centre-determined normal about equally well.
515    fn reference_body(
516        body_i: usize,
517        body_j: usize,
518        n_world: Vec3,
519        state: &State,
520        model: &Model,
521    ) -> usize {
522        // A sphere has no face and cannot own a normal; against a sphere the
523        // normal is the other body's. Deciding this by axis alignment instead
524        // ties at the identity (`+ẑ` is an axis of *every* upright frame), and
525        // the tie went to the ball: its frozen normal then rolled with it, and
526        // a marble at rest on a plate reported `dJ/dx₀ = 102` instead of `1`.
527        match (sphere_radius(model, body_i), sphere_radius(model, body_j)) {
528            (Some(_), None) => return body_j,
529            (None, Some(_)) => return body_i,
530            _ => {}
531        }
532        // How nearly a direction is a coordinate axis of a frame: 1 exactly on
533        // an axis, 1/sqrt(3) at the worst-case body diagonal.
534        let axis_alignment = |body: usize| -> f64 {
535            let local = state.body_xform[body].rot * n_world;
536            local.x.abs().max(local.y.abs()).max(local.z.abs())
537        };
538        if axis_alignment(body_j) > axis_alignment(body_i) {
539            body_j
540        } else {
541            body_i
542        }
543    }
544
545    /// Re-evaluate the collision this anchor stands for at (the FK of) a
546    /// perturbed state — same identity, smoothly moved geometry.
547    fn collision(&self, state: &State, ground_height: f64) -> Collision {
548        match *self {
549            Self::Ground {
550                body,
551                material_point,
552                world_offset,
553            } => {
554                let xform = &state.body_xform[body];
555                let support = xform.pos + xform.rot.transpose() * material_point + world_offset;
556                let depth = ground_height - support.z;
557                Collision {
558                    body_i: body,
559                    body_j: Collision::WORLD,
560                    contact_point: Vec3::new(support.x, support.y, ground_height - depth * 0.5),
561                    contact_normal: Vec3::z(),
562                    penetration_depth: depth,
563                }
564            }
565            Self::GroundCylinderRim {
566                body,
567                cap_centre,
568                axis_local,
569                radius,
570                quarter_turns,
571                frozen_offset,
572            } => {
573                let xform = &state.body_xform[body];
574                let axis = xform.rot.transpose() * axis_local;
575                // The same function detection used, so an anchor cannot drift
576                // from the detector that produced it.
577                let offset =
578                    cylinder_rim_dir(axis, quarter_turns).map_or(frozen_offset, |dir| dir * radius);
579                let support = xform.pos + xform.rot.transpose() * cap_centre + offset;
580                let depth = ground_height - support.z;
581                Collision {
582                    body_i: body,
583                    body_j: Collision::WORLD,
584                    contact_point: Vec3::new(support.x, support.y, ground_height - depth * 0.5),
585                    contact_normal: Vec3::z(),
586                    penetration_depth: depth,
587                }
588            }
589            Self::Pair {
590                body_i,
591                body_j,
592                point_i,
593                point_j,
594                normal_frame,
595                normal_local,
596                radius_i,
597                radius_j,
598            } => {
599                let xi = &state.body_xform[body_i];
600                let xj = &state.body_xform[body_j];
601                // The frozen body-frame normal stays unit under a rotation, so
602                // no renormalization is needed and none is done: a `normalize`
603                // here would divide by a quantity that is identically one, and
604                // its derivative would be a spurious zero-magnitude channel.
605                let n = state.body_xform[normal_frame].rot.transpose() * normal_local;
606                // `i` separates along `+n`, so `i` sits on the `+n` side and
607                // its surface point faces `−n`; the mirror image for `j`.
608                let mut pi = xi.pos + xi.rot.transpose() * point_i;
609                let mut pj = xj.pos + xj.rot.transpose() * point_j;
610                if let Some(r) = radius_i {
611                    pi -= n * r;
612                }
613                if let Some(r) = radius_j {
614                    pj += n * r;
615                }
616                // Positive = overlapping, consistent with the ground branch:
617                // the surfaces have swapped sides along `n` by this much.
618                let depth = (pj - pi).dot(n);
619                // The contact point rides the body that owns the *vertex*, not
620                // the midpoint of the two anchors.
621                //
622                // Averaging looks natural and is wrong, in a way that is only
623                // visible in a body-body pair. A face-vertex contact is located
624                // by the vertex; the face body contributes the plane the vertex
625                // is measured against, and sliding the face body *within its own
626                // plane* must not move the contact at all. `(pi + pj)/2` instead
627                // moves it by half of any such slide, which shows up as a
628                // spurious moment arm: the measured symptom was
629                // `dJ/d(plank x) = +9.7e-3` on a block sliding along a plank,
630                // where translation invariance puts the true value at `-8.2e-5`.
631                // The ground branch never had the bug because it already pins
632                // the point to `body_i` and takes only `z` from the plane —
633                // this is that same rule with "the plane" generalized to
634                // whichever body owns the normal.
635                let vertex = if normal_frame == body_i { pj } else { pi };
636                let sign = if normal_frame == body_i { -1.0 } else { 1.0 };
637                Collision {
638                    body_i,
639                    body_j,
640                    contact_point: vertex + n * (sign * 0.5 * depth),
641                    contact_normal: n,
642                    penetration_depth: depth,
643                }
644            }
645        }
646    }
647}
648
649// ---------------------------------------------------------------------------
650// Forward pass
651// ---------------------------------------------------------------------------
652
653/// Everything the backward pass needs about one step, cached rather than
654/// recomputed: re-solving during the backward pass would re-run the
655/// warm-started solver on a different warm-start history, which is
656/// path-dependent — the cache is what guarantees the adjoint differentiates
657/// the solves the forward pass actually performed.
658struct StepRecord {
659    q: DVec,
660    v: DVec,
661    u: DVec,
662    v_next: DVec,
663    anchors: Vec<Anchor>,
664    contact: Option<(ContactAssembly, ContactSolution)>,
665}
666
667/// The assembled problem and its solution for a contacted step (`None` when
668/// the step had no contacts).
669type StepSolve = Option<(ContactAssembly, ContactSolution)>;
670
671/// One step of the forward pass, mirroring `Simulator::step_with_contacts`
672/// operation for operation (FK → detect → free velocity → assemble →
673/// warm-started solve → integrate → FK).
674fn forward_step(
675    model: &Model,
676    state: &mut State,
677    ground_height: f64,
678    material: &ContactMaterial,
679    config: &ContactSolverConfig,
680    cache: &mut ContactCache,
681) -> (Vec<(Collision, GroundSupport)>, StepSolve) {
682    let dt = model.dt;
683
684    let (xforms, _velocities) = forward_kinematics(model, state);
685    state.body_xform = xforms;
686
687    // Same detection as `Simulator::step_with_contacts` — the full collision
688    // set — plus how each ground contact attaches to its body, which the
689    // anchor recovery below needs. Body-body contacts carry `Material`; it is
690    // never read, since they are rejected as unsupported first.
691    let mut contacts: Vec<(Collision, GroundSupport)> =
692        find_ground_contacts_model_with_support(model, state, ground_height, material.margin);
693    let body_contacts = find_contacts(model, state, material.margin);
694    contacts.extend(
695        body_contacts
696            .into_iter()
697            .map(|c| (c, GroundSupport::Material)),
698    );
699
700    // Mirrors `Simulator::step_with_contacts` bit for bit: a free joint's
701    // body-frame turn is stripped from `qdd` here and put back, exactly,
702    // after the solve (`phyz_rigid::strip_free_joint_coriolis`).
703    let mut qdd = aba(model, state);
704    let v_before = state.v.clone();
705    strip_free_joint_coriolis(model, v_before.as_slice(), qdd.as_mut_slice());
706    let free_qd = &state.v + &(&qdd * dt);
707
708    let bare: Vec<Collision> = contacts.iter().map(|(c, _)| c.clone()).collect();
709    let record = if contacts.is_empty() {
710        state.v = free_qd;
711        None
712    } else {
713        let materials = model.contact_materials(material);
714        let asm = assemble(model, state, &bare, &materials, &free_qd, dt, config);
715        let seed = cache.warm_start(state, &bare);
716        let solution = solve_contacts_warm(&asm.problem, config, &seed);
717        cache.store(state, &bare, &solution.impulses);
718        state.v = &free_qd + &asm.velocity_delta(&solution.impulses);
719        Some((asm, solution))
720    };
721    rotate_free_joint_velocities(model, v_before.as_slice(), state.v.as_mut_slice(), dt);
722
723    let v_clone = state.v.clone();
724    integrate_configuration(model, state.q.as_mut_slice(), v_clone.as_slice(), dt);
725    state.time += dt;
726
727    let (xforms, _) = forward_kinematics(model, state);
728    state.body_xform = xforms;
729
730    (contacts, record)
731}
732
733/// Run the nominal rollout, recording every step.
734fn forward_rollout(
735    rollout: &ConvexContactRollout,
736) -> Result<(Vec<StepRecord>, State), ConvexAdjointError> {
737    let model = rollout.model;
738    let mut state = model.default_state();
739    state.q = rollout.q0.clone();
740    state.v = rollout.v0.clone();
741    let mut cache = ContactCache::default();
742    let mut records = Vec::with_capacity(rollout.steps);
743
744    for t in 0..rollout.steps {
745        let q = state.q.clone();
746        let v = state.v.clone();
747        let u = (rollout.ctrl)(t);
748        state.ctrl = u.clone();
749
750        // Anchors are read off the *pre-step* FK, the same transforms the
751        // detection inside `forward_step` will use.
752        let (xforms, _) = forward_kinematics(model, &state);
753        state.body_xform = xforms;
754
755        let (contacts, contact) = forward_step(
756            model,
757            &mut state,
758            rollout.ground_height,
759            &rollout.material,
760            &rollout.config,
761            &mut cache,
762        );
763
764        let mut anchors = Vec::with_capacity(contacts.len());
765        // Recover anchors against the pre-step transforms; `forward_step`
766        // already advanced `state`, so rebuild them from (q, v).
767        let mut pre = model.default_state();
768        pre.q = q.clone();
769        pre.v = v.clone();
770        let (pre_xf, _) = forward_kinematics(model, &pre);
771        pre.body_xform = pre_xf;
772        for (c, support) in &contacts {
773            anchors.push(Anchor::of(c, *support, &pre, rollout.ground_height, model));
774        }
775
776        records.push(StepRecord {
777            q,
778            v,
779            u,
780            v_next: state.v.clone(),
781            anchors,
782            contact,
783        });
784    }
785    Ok((records, state))
786}
787
788/// Run the nominal rollout and return the objective only — the FD oracle for
789/// gates, and a cheap primal probe. The dynamics are identical to
790/// [`convex_adjoint_gradient`]'s forward pass (and to
791/// `Simulator::step_with_contacts`).
792pub fn convex_rollout_objective(
793    rollout: &ConvexContactRollout,
794    objective: &FinalStateObjective,
795) -> f64 {
796    let model = rollout.model;
797    let mut state = model.default_state();
798    state.q = rollout.q0.clone();
799    state.v = rollout.v0.clone();
800    let mut cache = ContactCache::default();
801    for t in 0..rollout.steps {
802        state.ctrl = (rollout.ctrl)(t);
803        forward_step(
804            model,
805            &mut state,
806            rollout.ground_height,
807            &rollout.material,
808            &rollout.config,
809            &mut cache,
810        );
811    }
812    (objective.value)(state.q.as_slice(), state.v.as_slice())
813}
814
815// ---------------------------------------------------------------------------
816// Backward pass
817// ---------------------------------------------------------------------------
818
819/// The smooth pieces of one step at a (possibly perturbed) input, with the
820/// contact set frozen to the nominal anchors and the impulses frozen to `f*`.
821///
822/// `residual` is the KKT residual `(A + R)f* + b − e_n·bias` — the object
823/// whose differential the IFT contracts. `gf` is `M⁻¹ Jᵀ f*`, the explicit
824/// (impulse-held-fixed) contact term of `v'`.
825struct Pieces {
826    v_free: DVec,
827    residual: Vec<f64>,
828    /// Per contact, `M_t(θ)·t_rel` for sliding contacts (`A`'s own tangential
829    /// 2x2 block applied to `t* − f_t`; the regularizer floor is constant and
830    /// drops out of the difference). Zero for non-sliding contacts.
831    mt_rel: Vec<[f64; 2]>,
832    gf: DVec,
833}
834
835// Ten arguments: this is the frozen step function itself — every input is a
836// genuine independent variable of the map being differentiated.
837#[allow(clippy::too_many_arguments)]
838// Stride arithmetic into the flat, row-major Delassus (base = 3*c).
839#[allow(clippy::needless_range_loop)]
840fn eval_pieces(
841    model: &Model,
842    ground_height: f64,
843    material: &ContactMaterial,
844    config: &ContactSolverConfig,
845    anchors: &[Anchor],
846    f_star: &[Vec3],
847    t_rel: &[[f64; 2]],
848    q: &DVec,
849    v: &DVec,
850    u: &DVec,
851) -> Pieces {
852    let dt = model.dt;
853    let mut state = model.default_state();
854    state.q = q.clone();
855    state.v = v.clone();
856    state.ctrl = u.clone();
857    let (xforms, _) = forward_kinematics(model, &state);
858    state.body_xform = xforms;
859
860    // The same stripped `qdd` the forward step integrates. The turn that
861    // follows the solve (`rotate_free_joint_velocities`) is linear in `v'`
862    // and is applied to each lane's `dv'` by `rotate_tangent` below.
863    let mut qdd = aba(model, &state);
864    strip_free_joint_coriolis(model, state.v.as_slice(), qdd.as_mut_slice());
865    let v_free = &state.v + &(&qdd * dt);
866
867    if anchors.is_empty() {
868        return Pieces {
869            v_free,
870            residual: Vec::new(),
871            mt_rel: Vec::new(),
872            gf: DVec::zeros(model.nv),
873        };
874    }
875
876    let contacts: Vec<Collision> = anchors
877        .iter()
878        .map(|a| a.collision(&state, ground_height))
879        .collect();
880    let materials = model.contact_materials(material);
881    let asm = assemble(model, &state, &contacts, &materials, &v_free, dt, config);
882    let p = &asm.problem;
883    let n = p.n;
884    let dim = 3 * n;
885    let mut flat = vec![0.0; dim];
886    for (c, f) in f_star.iter().enumerate() {
887        flat[3 * c] = f.x;
888        flat[3 * c + 1] = f.y;
889        flat[3 * c + 2] = f.z;
890    }
891    let mut residual = vec![0.0; dim];
892    for c in 0..n {
893        let reg = regularization_diag(p, c, config);
894        let base = 3 * c;
895        for r in 0..3 {
896            let mut acc = p.free_velocity[base + r];
897            for (col, fc) in flat.iter().enumerate() {
898                acc += p.delassus[(base + r) * dim + col] * fc;
899            }
900            acc += reg[r] * flat[base + r];
901            if r == 0 {
902                acc -= p.rows[c].bias;
903            }
904            residual[base + r] = acc;
905        }
906    }
907    let mut mt_rel = vec![[0.0f64; 2]; n];
908    for (c, tr) in t_rel.iter().enumerate() {
909        let base = 3 * c;
910        for i in 0..2 {
911            mt_rel[c][i] = p.delassus[(base + 1 + i) * dim + base + 1] * tr[0]
912                + p.delassus[(base + 1 + i) * dim + base + 2] * tr[1];
913        }
914    }
915    let gf = asm.velocity_delta(f_star);
916
917    Pieces {
918        v_free,
919        residual,
920        mt_rel,
921        gf,
922    }
923}
924
925/// Tangent of `rotate_free_joint_velocities` at this step.
926///
927/// The forward pass turns each free joint's linear velocity by `R(−ω dt)`
928/// after the solve, with `ω` read from the step's input velocity `v`. So
929/// `dv'' = R(−ω dt)·dv' + (∂R(−ω dt)/∂ω · dω)·v'`, where `v'` is the
930/// pre-turn velocity (recovered from the recorded post-turn `v_next` by the
931/// inverse turn) and `dω` is nonzero only on a `v` lane. The second term is
932/// `O(dt)`, the same order as everything else in the step's Jacobian, and
933/// is taken by central difference of the exact turn.
934fn rotate_tangent(model: &Model, v: &DVec, v_next: &DVec, dv: &mut DVec, v_lane: Option<usize>) {
935    let dt = model.dt;
936    for (jidx, joint) in model.joints.iter().enumerate() {
937        if joint.joint_type != phyz_model::JointType::Free {
938            continue;
939        }
940        let off = model.v_offsets[jidx];
941        let omega = Vec3::new(v[off], v[off + 1], v[off + 2]);
942        let d = Vec3::new(dv[off + 3], dv[off + 4], dv[off + 5]);
943        let mut out = phyz_math::quat_exp(&(omega * -dt)).rotate(d);
944        if let Some(j) = v_lane
945            && j >= off
946            && j < off + 3
947        {
948            let domega = match j - off {
949                0 => Vec3::new(1.0, 0.0, 0.0),
950                1 => Vec3::new(0.0, 1.0, 0.0),
951                _ => Vec3::new(0.0, 0.0, 1.0),
952            };
953            let after = Vec3::new(v_next[off + 3], v_next[off + 4], v_next[off + 5]);
954            let before = phyz_math::quat_exp(&(omega * dt)).rotate(after);
955            let h = 1e-6;
956            let p = phyz_math::quat_exp(&((omega + domega * h) * -dt)).rotate(before);
957            let m = phyz_math::quat_exp(&((omega - domega * h) * -dt)).rotate(before);
958            out += (p - m) * (1.0 / (2.0 * h));
959        }
960        dv[off + 3] = out.x;
961        dv[off + 4] = out.y;
962        dv[off + 5] = out.z;
963    }
964}
965
966/// `q' = Φ(q, v')` — the shared configuration update.
967fn phi(model: &Model, q: &DVec, v_next: &DVec) -> DVec {
968    let mut out = q.clone();
969    integrate_configuration(model, out.as_mut_slice(), v_next.as_slice(), model.dt);
970    out
971}
972
973/// A model with body `b`'s inertia rebuilt from a perturbed 10-vector.
974fn perturbed_model(model: &Model, body: usize, params: &[f64; N_INERTIA_PARAMS]) -> Model {
975    use phyz_math::{Mat3, SpatialInertia};
976    let mut m = model.clone();
977    let p = params;
978    m.bodies[body].inertia = SpatialInertia::new(
979        p[0],
980        Vec3::new(p[1], p[2], p[3]),
981        Mat3::new(p[4], p[7], p[8], p[7], p[5], p[9], p[8], p[9], p[6]),
982    );
983    m
984}
985
986/// One nominal rollout forward, one adjoint pass backward: `J`, `dJ/dq0`,
987/// `dJ/dv0`, `dJ/du_t`, and `dJ/dπ` per body — all through the convex contact
988/// model the forward simulator integrates.
989///
990/// Errors instead of returning a wrong gradient; see [`ConvexAdjointError`].
991pub fn convex_adjoint_gradient(
992    rollout: &ConvexContactRollout,
993    objective: &FinalStateObjective,
994) -> Result<ConvexAdjointGradients, ConvexAdjointError> {
995    let model = rollout.model;
996    let nq = model.nq;
997    let nv = model.nv;
998    let nb = model.bodies.len();
999    let (records, final_state) = forward_rollout(rollout)?;
1000
1001    let j0 = (objective.value)(final_state.q.as_slice(), final_state.v.as_slice());
1002    let (gq, gv) = (objective.gradient)(final_state.q.as_slice(), final_state.v.as_slice());
1003    assert_eq!(gq.len(), nq, "objective ∂g/∂q length");
1004    assert_eq!(gv.len(), nv, "objective ∂g/∂v length");
1005    let mut lam_q = DVec::from_slice(&gq);
1006    let mut lam_v = DVec::from_slice(&gv);
1007
1008    let mut d_ctrl = vec![DVec::zeros(nv); rollout.steps];
1009    let mut d_inertia = vec![[0.0f64; N_INERTIA_PARAMS]; nb];
1010    let mut d_friction = 0.0f64;
1011    let mut d_restitution = 0.0f64;
1012    let nominal_params: Vec<[f64; N_INERTIA_PARAMS]> = model
1013        .bodies
1014        .iter()
1015        .map(|b| inertia_params(&b.inertia))
1016        .collect();
1017
1018    for (t, rec) in records.iter().enumerate().rev() {
1019        // IFT sensitivity of this step's *recorded* solve. Cached, never
1020        // re-solved: warm starting makes a re-solve path-dependent.
1021        let fps: Option<FixedPointSensitivity> = match &rec.contact {
1022            None => None,
1023            Some((asm, sol)) => {
1024                if !sol.converged {
1025                    return Err(ConvexAdjointError::Unconverged {
1026                        step: t,
1027                        iterations: sol.iterations,
1028                        residual: sol.residual,
1029                    });
1030                }
1031                Some(
1032                    FixedPointSensitivity::at(&asm.problem, sol, &rollout.config)
1033                        .ok_or(ConvexAdjointError::SingularKkt { step: t })?,
1034                )
1035            }
1036        };
1037        let t_rel: Vec<[f64; 2]> = match (&rec.contact, &fps) {
1038            (Some((asm, _)), Some(s)) => (0..asm.problem.n)
1039                .map(|c| s.slide_tangent(c).map(|st| st.t_rel).unwrap_or([0.0; 2]))
1040                .collect(),
1041            _ => Vec::new(),
1042        };
1043        let f_star: Vec<Vec3> = rec
1044            .contact
1045            .as_ref()
1046            .map(|(_, s)| s.impulses.clone())
1047            .unwrap_or_default();
1048
1049        // Directional derivative of v' from a directional derivative of the
1050        // smooth pieces, closing the contact channel with the IFT.
1051        let dv_next = |dp: &Pieces, v_lane: Option<usize>| -> DVec {
1052            let mut dvn = match (&rec.contact, &fps) {
1053                (Some((asm, _)), Some(s)) => {
1054                    // df* from the exact map linearization; dv' = dv_free +
1055                    // d(M⁻¹Jᵀ f*)|_f + M⁻¹Jᵀ df*.
1056                    let df = s.apply(&dp.residual, &dp.mt_rel);
1057                    &(&dp.v_free + &dp.gf) + &asm.velocity_delta(&df)
1058                }
1059                _ => dp.v_free.clone(),
1060            };
1061            // ...then the free joints' turn into the end-of-step frame.
1062            rotate_tangent(model, &rec.v, &rec.v_next, &mut dvn, v_lane);
1063            dvn
1064        };
1065
1066        let eval = |m: &Model, mat: &ContactMaterial, q: &DVec, v: &DVec, u: &DVec| -> Pieces {
1067            eval_pieces(
1068                m,
1069                rollout.ground_height,
1070                mat,
1071                &rollout.config,
1072                &rec.anchors,
1073                &f_star,
1074                &t_rel,
1075                q,
1076                v,
1077                u,
1078            )
1079        };
1080
1081        // Central difference of the pieces along one input lane.
1082        //
1083        // The material is a lane input like the model is: restitution reaches
1084        // the impulses through `b`, so perturbing it perturbs the residual and
1085        // this machinery prices it with no special case. Friction does *not*
1086        // appear in the residual at all — it lives only in the cone
1087        // constraint — so it cannot be a lane and is handled separately below.
1088        #[allow(clippy::too_many_arguments)]
1089        let lane = |mp: &Model,
1090                    mm: &Model,
1091                    matp: &ContactMaterial,
1092                    matm: &ContactMaterial,
1093                    qp: &DVec,
1094                    qm: &DVec,
1095                    vp: &DVec,
1096                    vm: &DVec,
1097                    up: &DVec,
1098                    um: &DVec,
1099                    h: f64|
1100         -> Pieces {
1101            let a = eval(mp, matp, qp, vp, up);
1102            let b = eval(mm, matm, qm, vm, um);
1103            let inv = 1.0 / (2.0 * h);
1104            Pieces {
1105                v_free: &(&a.v_free - &b.v_free) * inv,
1106                residual: a
1107                    .residual
1108                    .iter()
1109                    .zip(&b.residual)
1110                    .map(|(x, y)| (x - y) * inv)
1111                    .collect(),
1112                mt_rel: a
1113                    .mt_rel
1114                    .iter()
1115                    .zip(&b.mt_rel)
1116                    .map(|(x, y)| [(x[0] - y[0]) * inv, (x[1] - y[1]) * inv])
1117                    .collect(),
1118                gf: &(&a.gf - &b.gf) * inv,
1119            }
1120        };
1121
1122        // dq' along a v'-direction, by central difference through Φ.
1123        let dphi_dvnext = |dvn: &DVec| -> DVec {
1124            let norm = dvn.as_slice().iter().fold(0.0f64, |m, x| m.max(x.abs()));
1125            if norm == 0.0 {
1126                return DVec::zeros(nq);
1127            }
1128            let s = FD_EPS / norm;
1129            let vp = &rec.v_next + &(dvn * s);
1130            let vm = &rec.v_next - &(dvn * s);
1131            &(&phi(model, &rec.q, &vp) - &phi(model, &rec.q, &vm)) * (1.0 / (2.0 * s))
1132        };
1133
1134        // Contract a lane's (dq', dv') against the incoming adjoint.
1135        let contract = |dqn: &DVec, dvn: &DVec| -> f64 {
1136            let mut acc = 0.0;
1137            for i in 0..nq {
1138                acc += lam_q[i] * dqn[i];
1139            }
1140            for j in 0..nv {
1141                acc += lam_v[j] * dvn[j];
1142            }
1143            acc
1144        };
1145
1146        let mut new_lam_q = DVec::zeros(nq);
1147        let mut new_lam_v = DVec::zeros(nv);
1148
1149        // --- q lanes ---
1150        for i in 0..nq {
1151            let h = FD_EPS * rec.q[i].abs().max(1.0);
1152            let mut qp = rec.q.clone();
1153            let mut qm = rec.q.clone();
1154            qp[i] += h;
1155            qm[i] -= h;
1156            let dp = lane(
1157                model,
1158                model,
1159                &rollout.material,
1160                &rollout.material,
1161                &qp,
1162                &qm,
1163                &rec.v,
1164                &rec.v,
1165                &rec.u,
1166                &rec.u,
1167                h,
1168            );
1169            let dvn = dv_next(&dp, None);
1170            // Direct Φ_q block plus the v'-mediated part.
1171            let dqn_direct = &(&phi(model, &qp, &rec.v_next) - &phi(model, &qm, &rec.v_next))
1172                * (1.0 / (2.0 * h));
1173            let dqn = &dqn_direct + &dphi_dvnext(&dvn);
1174            new_lam_q[i] = contract(&dqn, &dvn);
1175        }
1176
1177        // --- v lanes ---
1178        for j in 0..nv {
1179            let h = FD_EPS * rec.v[j].abs().max(1.0);
1180            let mut vp = rec.v.clone();
1181            let mut vm = rec.v.clone();
1182            vp[j] += h;
1183            vm[j] -= h;
1184            let dp = lane(
1185                model,
1186                model,
1187                &rollout.material,
1188                &rollout.material,
1189                &rec.q,
1190                &rec.q,
1191                &vp,
1192                &vm,
1193                &rec.u,
1194                &rec.u,
1195                h,
1196            );
1197            let dvn = dv_next(&dp, Some(j));
1198            let dqn = dphi_dvnext(&dvn);
1199            new_lam_v[j] = contract(&dqn, &dvn);
1200        }
1201
1202        // --- control lanes ---
1203        for j in 0..nv {
1204            let h = FD_EPS * rec.u[j].abs().max(1.0);
1205            let mut up = rec.u.clone();
1206            let mut um = rec.u.clone();
1207            up[j] += h;
1208            um[j] -= h;
1209            let dp = lane(
1210                model,
1211                model,
1212                &rollout.material,
1213                &rollout.material,
1214                &rec.q,
1215                &rec.q,
1216                &rec.v,
1217                &rec.v,
1218                &up,
1219                &um,
1220                h,
1221            );
1222            let dvn = dv_next(&dp, None);
1223            let dqn = dphi_dvnext(&dvn);
1224            d_ctrl[t][j] = contract(&dqn, &dvn);
1225        }
1226
1227        // --- inertia-parameter lanes ---
1228        for b in 0..nb {
1229            for k in 0..N_INERTIA_PARAMS {
1230                let h = FD_EPS * nominal_params[b][k].abs().max(1.0);
1231                let mut pp = nominal_params[b];
1232                let mut pm = nominal_params[b];
1233                pp[k] += h;
1234                pm[k] -= h;
1235                let mp = perturbed_model(model, b, &pp);
1236                let mm = perturbed_model(model, b, &pm);
1237                let dp = lane(
1238                    &mp,
1239                    &mm,
1240                    &rollout.material,
1241                    &rollout.material,
1242                    &rec.q,
1243                    &rec.q,
1244                    &rec.v,
1245                    &rec.v,
1246                    &rec.u,
1247                    &rec.u,
1248                    h,
1249                );
1250                let dvn = dv_next(&dp, None);
1251                let dqn = dphi_dvnext(&dvn);
1252                d_inertia[b][k] += contract(&dqn, &dvn);
1253            }
1254        }
1255
1256        // --- restitution lane ---
1257        //
1258        // `e` reaches the impulses through `b`: assembly scales the normal row
1259        // of the free velocity by `1 + e_eff`. So it is an ordinary lane, and
1260        // the machinery above prices it including the low-speed smoothstep
1261        // ramp, which is exactly why §4.3 insisted restitution be a term in `b`
1262        // rather than a post-solve velocity reset — a reset would be a branch
1263        // on the primal, with no derivative in `e` at all at `v_n = 0`.
1264        {
1265            let h = FD_EPS * rollout.material.restitution.abs().max(1.0);
1266            let mut matp = rollout.material.clone();
1267            let mut matm = rollout.material.clone();
1268            matp.restitution += h;
1269            matm.restitution -= h;
1270            let dp = lane(
1271                model, model, &matp, &matm, &rec.q, &rec.q, &rec.v, &rec.v, &rec.u, &rec.u, h,
1272            );
1273            let dvn = dv_next(&dp, None);
1274            let dqn = dphi_dvnext(&dvn);
1275            d_restitution += contract(&dqn, &dvn);
1276        }
1277
1278        // --- friction channel ---
1279        //
1280        // Not a lane, because `mu` does not appear in the residual
1281        // `(A + R)f* + b − e_n·bias` anywhere: the friction coefficient enters
1282        // only through the *cone constraint*. Differencing `eval_pieces` in
1283        // `mu` would therefore return an exact, entirely convincing zero.
1284        //
1285        // `friction_sensitivity` supplies `df*/dmu` directly, already routed
1286        // through the coupled system (a change in one sliding contact's
1287        // tangential capacity moves every other contact through `A`). Only
1288        // sliding contacts have a non-zero column. Summing the columns gives
1289        // the derivative with respect to the single shared coefficient, and
1290        // from there the contraction is the same as every other lane's, since
1291        // `v_free` and the impulse-held-fixed term do not depend on `mu`.
1292        if let Some((asm, sol)) = &rec.contact {
1293            let n = asm.problem.n;
1294            if let Some(dfdmu) = friction_sensitivity(&asm.problem, sol, &rollout.config) {
1295                let mut df = vec![Vec3::zeros(); n];
1296                for (c, dfc) in df.iter_mut().enumerate() {
1297                    let base = 3 * c;
1298                    // Sum over columns: every contact shares one `mu`.
1299                    let mut acc = [0.0f64; 3];
1300                    for (r, a) in acc.iter_mut().enumerate() {
1301                        for col in 0..n {
1302                            *a += dfdmu[(base + r) * n + col];
1303                        }
1304                    }
1305                    *dfc = Vec3::new(acc[0], acc[1], acc[2]);
1306                }
1307                let dvn = asm.velocity_delta(&df);
1308                let dqn = dphi_dvnext(&dvn);
1309                d_friction += contract(&dqn, &dvn);
1310            }
1311        }
1312
1313        lam_q = new_lam_q;
1314        lam_v = new_lam_v;
1315    }
1316
1317    Ok(ConvexAdjointGradients {
1318        objective: j0,
1319        d_q0: lam_q,
1320        d_v0: lam_v,
1321        d_ctrl,
1322        d_inertia,
1323        d_friction,
1324        d_restitution,
1325    })
1326}