JAXSIM

jaxsim is a JAX rewrite of gradSim: differentiable rendering and physics for parameter estimation and control from video. Rigid bodies, deformable solids, cloth. All differentiable through simulation and rendering.

gradSim was PyTorch + NVIDIA warp/dflex. jaxsim replaces that with JAX: autodiff, JIT and physics integration, all native, no CUDA kernel compilation.

Comparing this implementation to the original PyTorch based codebase.

DFlex spring-mass chain (jit-able functional state)

n/a means the unrolled loop was too large to compile in reasonable time. Eager mode (no jit) runs about 150 ms flat, which is 20-70x slower than torch, and compilation itself costs 150-440 ms.

RigidBody Simulator (mutates state each step and not jit-able)

torch has ~10x faster throughout. Neither side JIT-compiles here because JIT compilation mutates Python attributes each step.

Optimization loop (30 Adam iterations, 50-particle chain, 20 sim-steps/iter)

Breakeven is about 690 iterations, below which the total wall-clock favors torch (compile cost dominates), and above which jax is faster.

DFlex spring-mass chain

Both stay flat across particle count here, because this problem size is kernel-launch bound rather than compute bound.

RigidBody Simulator

torch is about 8-9x faster on the forward pass and about 4x faster on the gradient.

Optimization loop (same task as CPU)

Breakeven is about 770 iterations, below which the total wall-clock favors torch, and above which jax is faster.

torch and jax were both slower on GPU than on CPU across every benchmark above, because these problem sizes (tens to low hundreds of particles or bodies) are too small to hide GPU launch and dispatch latency behind compute. Larger batches or bigger meshes would be needed to see the GPU actually pay off.

Requires Python ≥ 3.9.

python -m venv .venv && source .venv/bin/activate

pip install jax jaxlib

pip install -U "jax[cuda12]"

pip install -e .

Run any example from the repo root:

cd examples
../../jaxenv/bin/python3 <script>.py