Ostrich: Taking Large Strides Through Stiff Contact in Differentiable Dynamics

AI summary

Overview

• Research area: Robotics, differentiable physics simulation, hard-contact dynamics, and gradient-based trajectory optimization.
• Technical level: Advanced. The paper assumes familiarity with complementarity problems, Newton solvers, the Implicit Function Theorem, Schur complements, and GPU-parallel simulation.
• Scope: Ostrich is a GPU-accelerated differentiable rigid-body simulator that resolves hard contact and friction with a non-smooth Newton solver at large timesteps and backpropagates through the converged residual at constant memory per step, enabling gradient-based control optimization through stiff contact and triangle-mesh terrain.

What This Paper Is About

Differentiable simulators let robots learn and optimize control policies through gradients, but on tasks with stiff, persistent contact — a wheeled robot climbing a pallet, for example — existing engines force a three-way tradeoff between simulation accuracy, gradient reliability, and per-iteration cost. Tape-based engines such as MJX need very small timesteps to keep smoothed contacts numerically tractable, and their memory grows linearly with the number of timesteps; surrogate models bound memory but discard the contact geometry the optimizer needs. Ostrich's goal is to break that tradeoff by combining hard contacts, a large-timestep implicit integrator, and an adjoint backward pass whose memory is independent of the inner solver's iteration count.

Key Contributions

1. A GPU-accelerated differentiable simulator with hard NCP contact and Coulomb friction, resolved by an inexact non-smooth Newton solver inside each backward-Euler step, stable at macroscopic timesteps on the order of 10⁻¹ s.
2. An adjoint backward pass that reuses the forward solve's Schur complement and preconditioner. One PCR solve on the converged residual yields reverse-mode gradients at O(1) memory per timestep, decoupling gradient memory from the number of inner contact iterations.
3. Solver extensions beyond the non-smooth Newton formulation of Macklin et al.: implicit control/servo constraints solved jointly with contact (well posed at any positive gain), best-iterate backtracking, and a robust friction projection that keeps intermediate iterates cone-admissible.
4. Scaling to arbitrary triangle meshes and thousands of parallel worlds on a single GPU, demonstrated with a 10 s gradient-based terrain-traversal optimization over a ~12,500-face mesh and 8,192 parallel worlds differentiated on one 24 GB card.

Main Findings

• Sim-to-real accuracy at 50× larger timestep. On 14 real traversals of a 14.4 cm pallet (three-wheeled, 106 kg robot), Ostrich reaches a held-out combined position-and-yaw error of 0.208 m, compared with 0.268 m for Newton Semi-Implicit and 0.315 m for CPU MuJoCo. Ostrich holds its accuracy floor up to h = 0.1 s, roughly 50× larger than MuJoCo's usable step (2 ms) and 200× larger than Semi-Implicit's (0.5 ms). The residual error is attributed to a shared constant-friction model, not to any solver.

• Gradients converge where baselines fail. On a box-obstacle control-synthesis task (50 Adam iterations, 25 random trials), Ostrich succeeds on 100% of trials with a median final position error of 0.065 m. MJX succeeds on 16% (0.32 m median error) and characteristically approaches the target without stopping. Newton Semi-Implicit produces non-finite gradients beyond a ~2 s horizon and succeeds on only 8% of trials (0.58 m median error).

• Large per-iteration speedups. A warm optimization iteration takes 0.55 s for Ostrich versus 116 s for MJX (211× faster) and 2.6 s for Semi-Implicit (4.7× faster). End-to-end, a 50-iteration run costs 29 s for Ostrich versus 1.6 h for MJX and 0.5 h for Semi-Implicit, and neither baseline reaches a comparable loss.

• Parallel-world scaling. Ostrich differentiates 8,192 parallel worlds on a single 24 GB GPU, peaking at 348 world-iterations/s. Checkpointed MJX peaks at 12 world-iterations/s at its 4,096-world memory cap (a 29× gap), and Semi-Implicit peaks at 23 at 512 worlds. With checkpointing, Ostrich sustains 171–173 world-iterations/s up to 32,768 worlds on 21.7 GB.

• Efficient zero-order comparison. Evolution strategies match first-order optimization in wall-clock on a single problem (0.8× the budget of 50 Adam iterations), but spend about 1,800 rollouts per solved problem versus 50 forward-adjoint pairs (a 35× sample-efficiency gap) and degrade at 300 parameters.

• Mesh-terrain trajectory optimization. Across 50 random triangle-mesh instances (10 s horizon, T = 125, h = 80 ms), median best RMSE is 0.143 m (mean 0.208 ± 0.178 m). The full position pull-back amplifies contact sensitivity (~1.13× growth per step) over long horizons, so contact geometry is held fixed within each step.

• Adjoint verified numerically. Against centered finite differences, the adjoint shows 0.02% maximum relative error on an impulsive contact-boundary test, ≤0.4% across all wheel degrees of freedom, and ≤1.6% on cart-pole losses over horizons up to 50 steps, with checks running in continuous integration.

Methodology in Plain English

Ostrich represents the robot and its environment as rigid bodies in maximal coordinates and discretizes motion with backward Euler integration. Every physical interaction — joints, normal contact, friction, and motor servos — is written as a residual equation that must equal zero. Non-penetration is encoded with the Fischer-Burmeister function, a standard trick that turns a "either no gap or no force" condition into a smooth-looking equation with a root exactly at the complementarity solution. Friction is derived from the principle of maximal dissipation and projected onto the Coulomb cone, with a weight that interpolates between sticking and sliding; the weight is frozen inside each Newton iteration so the friction block becomes symmetric, and at convergence it matches the equilibrium value.

The nonlinear system is solved by an inexact non-smooth Newton method. Because the position update is linear in velocity, the position unknowns are eliminated by substitution, leaving a KKT block system in velocity and multipliers. A Schur complement then eliminates velocity, producing a smaller system solved matrix-free with Preconditioned Conjugate Residual iteration. Best-iterate backtracking returns the lowest-residual iterate from a fixed Newton budget rather than accepting the final one, and friction safeguards keep every returned iterate within the cone.

The key methodological move is on the backward pass. Instead of taping every inner solver iteration — which is what makes MJX and similar engines memory-hungry — Ostrich applies the Implicit Function Theorem to the converged residual. The adjoint system shares the exact same Schur complement matrix and preconditioner as the forward solve, so no additional matrix assembly is needed. One extra PCR solve yields all gradients. A single residual evaluation at the converged state is recorded on an AD tape so that the dependence of contact points and normals on the previous pose is captured automatically. Ostrich is implemented in NVIDIA Warp GPU kernels, using NVIDIA Newton for model building and collision detection.

Why This Matters

This work directly attacks the bottleneck that limits gradient-based robotics: contact. Without reliable, affordable gradients through stiff contact, differentiable simulators remain useful for smooth manipulation tasks but not for driving over rough terrain, climbing obstacles, or any setting where impact and friction dominate. Ostrich shows that hard contacts, large timesteps, and cheap gradients can coexist on a single GPU, which changes what long-horizon trajectory optimization and policy learning are practically feasible.

Real-world applications:

• Off-road and warehouse vehicle control: optimizing wheel-command splines through pallets, curbs, and uneven ground for autonomous forklifts and delivery robots.
• Legged and wheeled robot locomotion: training or refining controllers where foot or wheel impacts are stiff and long-horizon credit assignment matters.
• Model-predictive control on GPU: interactive-rate single-world optimization (0.53 s per iteration) that is fast enough to close the loop for complex articulated platforms.
• Fleet-scale policy learning in simulation: differentiating thousands of parallel environments on one GPU for batched reinforcement learning and system identification.

Industry relevance: Robotics companies, autonomous-vehicle teams, and simulator vendors (MuJoCo, NVIDIA Isaac/Newton, Genesis) all depend on contact handling. Ostrich builds on NVIDIA's Warp and Newton stack, so its techniques are directly portable to existing GPU simulation pipelines, and its 24 GB single-GPU footprint makes high-fidelity differentiable contact accessible without a data-center cluster.

Future Directions

• Friction model fidelity. All three engines saturate at a similar error floor on turn-heavy real runs because a constant friction coefficient cannot capture the speed-dependent turn efficiency variation (0.11–0.30) observed in the data. Learning or state-dependent friction models are a natural next step.
• Long-horizon contact sensitivity. The full position pull-back amplifies contact sensitivity by roughly 1.13× per step, which required truncating it on the 10 s mesh horizon. Stabilizing or regularizing long-horizon adjoints through deforming contact geometry is an open problem.
• Closed-loop deployment. The paper demonstrates open-loop trajectory optimization and, separately, differentiable MPC-style iteration timing. Closing the loop with Ostrich inside a real controller — beyond the open-loop wheel-velocity spline demo — remains to be shown.
• Extending to the excluded engines' regimes. Genesis, Dojo, and TinyDiffSim were excluded for gradient, hardware, or geometry reasons. Re-benchmarking as those engines evolve, and testing Ostrich on deformables, cloth, or granular contact beyond rigid bodies, would clarify where the hard-contact formulation generalizes.

Target Audience

This paper is aimed at researchers and engineers working on differentiable simulation, contact dynamics, and gradient-based robot control. The most direct beneficiaries are graduate students and industry practitioners building trajectory optimizers or policy-learning pipelines that must handle stiff, persistent contact, as well as simulator developers interested in implicit differentiation, Schur-complement adjoints, and GPU-parallel complementarity solvers. Readers without background in numerical optimization, complementarity problems, or GPU simulation will find the theory section dense, though the experimental comparisons against MJX and Newton Semi-Implicit are accessible on their own.

Authors’ abstract

Three properties determine whether a differentiable simulator can drive gradient-based optimization through contact: simulation accuracy, gradient reliability, and per-iteration cost. Tape-based engines such as MJX and Newton Semi-Implicit require timesteps small enough to keep contacts numerically tractable, and their backpropagation memory grows linearly with the number of timesteps T. Surrogate models bound memory by approximating contact away, but the resulting gradients lose the geometry the optimization depends on. We present Ostrich, a GPU-accelerated rigid-body simulator that resolves hard contacts and friction with non-smooth Newton iteration at large timesteps (h ~ 0.1 s), and differentiates the converged residual via the implicit function theorem, reusing the forward Schur complement to compute the adjoint at O(1) memory per timestep. On real-robot trajectories over a pallet obstacle, Ostrich holds MuJoCo's sim-to-real accuracy up to a 50x larger timestep. Its gradients converge from random initializations where MJX descends slowly and Newton Semi-Implicit stalls; a warm iteration runs 211x faster than MJX's and 4.7x faster than Semi-Implicit's. On the same scene Ostrich differentiates 8,192 parallel worlds on a single 24 GB GPU, sustaining 29x checkpointed MJX's optimization throughput; without checkpointing both baselines exhaust memory at far fewer worlds. We close with a gradient-based trajectory optimization demonstration over triangle-mesh terrain across a 10 s horizon, a setting where prior engines either restrict to primitive geometry or face the convergence and memory limits shown above.