CppMjStep: Extended Custom Autograd Function for Differentiable MuJoCo Dynamics

I built CppMjStep to make MuJoCo differentiable inside PyTorch without slowing everything down. It is a C++ autograd extension that lets me backpropagate through entire simulation rollouts, supports batched runs, and includes sensor gradients so models that rely on IMU or force data can train end to end. The core is multithreaded C++ for speed and it is pip installable. In my benchmarks, large models with big batch sizes see more than a 2x speedup compared with a Python approach. I published it as a pip installable library.

Using it feels like a regular PyTorch layer. You pass a MuJoCo model and data, pick the number of steps to unroll, and call the module to step the dynamics while keeping everything in the computational graph. The goal is simple, make gradient based control and system identification in MuJoCo fast enough for real experiments.

The GitHub repo is available here.