CppMjStep: Extended Custom Autograd Function for Differentiable MuJoCo Dynamics

Description

An efficient integration between PyTorch and MuJoCo, implemented in C++. Enables automatic differentiation through MuJoCo simulation trajectories, allowing for gradient-based optimization of control policies directly within PyTorch. This implementation extends to sensors in the model, and uses C++ with multithreading to improve on speed (>2x speedup with large batch sizes for large models).

Features

pip Installable: Simple installation through pip. Refer INSTALL.md.

Efficient Gradient Computations: Significantly more efficient than naive Jacobian finite differencing calculations as it utilizes the built-in finite difference method in MuJoCo mjd_transitionFD.

Multi-Step Calculations: Provides the ability to estimate gradients over multiple simulation steps, by propagating gradients through the entire trajectory.

Batch Simulation Support: Enables batched simulations and gradient computations, significantly improving computational efficiency for large-scale experiments.

Sensor Support: Enables addition of sensors to the model, and allows backprop of sensor data through simulation with next to no cost to computation.

Comparison of execution times

Execution time for models with no sensors

Execution time for models with no sensors (Python v/s C++)

Execution time for models with sensors

Execution time for models with no sensors (Python v/s C++)

Usage

```python import torch import mujoco as mj from CppMjStep import MjStep, PyMjStep

Initialize MuJoCo model and data

xmlpath = 'path/to/your/model.xml' mjmodel = mj.MjModel.fromxmlpath(filename=xmlpath) mjdata = mj.MjData(mj_model)

Define MjStep layer

torchwrappedmodel = MjStep(mjmodel, mjdata, n_steps=5, multithread=False)

Define initial state and control input tensors

state = torch.rand(mjmodel.nq + mjmodel.nv + mjmodel.na + mjmodel.nsensordata, requiresgrad=True) ctrl = torch.rand(mjmodel.nu, requires_grad=True)

Compute next state

next_state = MjStep(state, ctrl) ```

Notes

• As of MuJoCo 3.1.2 the initial state passed to rollout() must include a time-step, such that nstate = mj_stateSize(model, mjtState.mjSTATE_FULLPHYSICS).
• Gradients are implicitly computed, to expose them refer MjStep/nn_wrapper.py
• Multithreading support has been added for differentiation. It is up to the user to pass the multithread flag at initialization. See MjStep/nn_wrapper.py for more.

Citation

If you use this package in your research, a citation would be appreciated:

@software{CppMjStep2025, author = {Chariar, Mukundan}, title = {{CppMjStep: Extended Custom Autograd Extension for Differentiable MuJoCo Dynamics}}, year = {2025}, version = {1.0}, howpublished = {\url{https://github.com/mukundan-chariar1/CppMjStep}}, }

Special thanks to Elad Sharony, whose DiffMjStep inspired this module. Portions of this code are derived from DiffMjStep, which is MIT licensed.