Ostrich: Taking Large Strides Through Stiff Contact in Differentiable Dynamics
🦩 Ostrich: Making Gradient-Based Robotics Work at Real-World Timesteps
As ML researchers and roboticists increasingly push simulated training environments towards reality, we often hit a fundamental wall: the physics simulation must be accurate enough to trust, but computationally efficient enough to run large optimization loops. When dealing with hard contacts and friction (like rolling over an obstacle or interacting with mesh terrain), standard differentiable dynamics simulators struggle.
The breakthrough arrives with Ostrich.
### ⚙️ The Problem: Physics vs. Computation
To train robots using gradient-based optimization in simulation, you need a system that accurately models physics (the forward pass) and can reliably compute the gradients (the backward pass). Historically, simulators like MJX or Newton Semi-Implicit faced severe limitations when tackling challenging contact scenarios:
- Small Timesteps: They required extremely small time steps ($ ext{h}$), making long simulations computationally prohibitive.
- Memory Explosion: The backpropagation memory usage grew linearly with the number of timesteps ($T$), quickly exceeding available GPU memory for complex tasks or large batch sizes.
- Accuracy Loss (Surrogates): Some solutions used surrogate models, but these often lost the crucial geometric fidelity needed for high-precision robotic control.
### 🤯 Ostrich’s Solution: Large Strides and Deep Differentiation
Ostrich fundamentally changes how we model dynamic contact. It is a GPU-accelerated rigid-body simulator that achieves several critical feats:
- Large Timesteps: It resolves hard contacts and friction using non-smooth Newton iteration at massive time steps (down to $h ext{~} 0.1 ext{ s}$). This means simulating real-world dynamics over longer periods with fewer computations.
- Efficient Backpropagation: Crucially, it differentiates the converged residual via the implicit function theorem. By cleverly reusing the forward Schur complement calculation for the adjoint, Ostrich computes backpropagating gradients while maintaining optimal $O(1)$ memory usage per timestep.
### 🚀 Benchmark Results That Change the Game
The results demonstrate a massive leap in capability over established simulators like MuJoCo and MJX:
✅ Accuracy: On real-robot trajectories (e.g., navigating a pallet), Ostrich maintains MuJoCo’s sim-to-real accuracy even up to a 50x larger timestep.
✅ Convergence & Speed: Where baselines like MJX struggled with slow descent or failed to converge, Ostrich provided stable gradient computation from random initializations. A warm iteration is shown to be $211 ext{x}$ faster than MJX and $4.7 ext{x}$ faster than Semi-Implicit.
✅ Scalability (The Memory Win): When simulating 8,192 parallel worlds on a single 24 GB GPU, Ostrich sustained an optimization throughput $29 ext{x}$ higher than checkpointed MJX—and unlike the baselines, it did not exhaust memory.
### ✨ The Future: Mesh and Long Horizons
Ostrich doesn’t stop at primitives. The authors close by demonstrating successful gradient-based optimization over complex triangle-mesh terrain across a $10 ext{ s}$ horizon—a setting where previous engines were limited to simple shapes or suffered from the convergence/memory constraints described above.
### 🛠️ Why This Matters for ML and Robotics
This research dramatically lowers the barrier to entry for deploying complex gradient-based learning pipelines (like Model Predictive Control or Reinforcement Learning) on highly accurate, physically challenging simulation environments. It enables researchers to train digital twins that are truly representative of real-world physics.
Read the full paper here: Ostrich: Taking Large Strides Through Stiff Contact in Differentiable Dynamics