Revive DiffMjStep as a lean PyTorch–MuJoCo autograd package

.gitignore

@@ -0,0 +1,8 @@
+__pycache__/
+*.py[cod]
+.venv/
+*.egg-info/
+dist/
+build/
+MUJOCO_LOG.TXT
+.idea/

CITATION.cff

@@ -1,7 +1,11 @@
-@software{DiffMjStep2024,
-  author = {Sharony, Elad},
-  title = {{DiffMjStep: Custom Autograd Extension for Differentiable MuJoCo Dynamics}},
-  year = {2024},
-  version = {1.0},
-  howpublished = {\url{https://github.com/EladSharony/DiffMjStep}},
-}
+cff-version: 1.2.0
+message: "If you use DiffMjStep, please cite this software."
+type: software
+title: "DiffMjStep: PyTorch Autograd for MuJoCo Rollouts"
+version: "0.2.0"
+date-released: "2026-06-28"
+authors:
+  - family-names: Sharony
+    given-names: Elad
+repository-code: "https://github.com/EladSharony/DiffMjStep"
+license: MIT

README.md

@@ -1,61 +1,174 @@
+# DiffMjStep
 
-# DiffMjStep: Custom Autograd Function for Differentiable MuJoCo Dynamics
+**PyTorch autograd for MuJoCo rollouts on CPU.**
 
-[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+<p align="center">
+  <img src="assets/benchmark.svg" width="760"
+       alt="DiffMjStep is 51x faster forward and 13x faster backward than a pure-Python reference of the same algorithm (batch 4096, MuJoCo 3.10, Intel i9-14900K).">
+</p>
 
-## Description
+**51× faster forward and 13× faster backward** than a pure-Python reference of the same
+algorithm (batch 4096, CPU) — MuJoCo transition-FD gradients without leaving PyTorch.
 
-An efficient integration between PyTorch and MuJoCo. 
-Enables automatic differentiation through MuJoCo simulation trajectories, 
-allowing for gradient-based optimization of control policies directly within PyTorch. 
+The forward pass executes MuJoCo. The custom backward applies MuJoCo transition
+Jacobians from `mjd_transitionFD` as a reverse-time VJP, so losses can differentiate with
+respect to the initial state and control sequence without moving the simulation to another
+framework.
 
-## Features
+```python
+x1 = mj_step(model, x0, u)
+loss = x1.square().sum()
+loss.backward()  # x0.grad and u.grad
+```
+
+## Why it exists
 
-**<img src="https://cdn-icons-png.flaticon.com/128/4285/4285622.png" width="32" height="32"> 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](https://mujoco.readthedocs.io/en/stable/APIreference/APIfunctions.html#mjd-transitionfd).
+- `loss.backward()` works through state, control, and optional sensor trajectories.
+- `backend="auto"` uses MuJoCo's C rollout and the native threaded backward when available,
+  with a clear Python fallback.
+- The Python fallback and native path implement the same MuJoCo transition-FD VJP.
+- The implementation stays small: Python orchestration plus one packaged C++ source file.
 
-**<img src="https://cdn-icons-png.flaticon.com/128/9072/9072147.png" width="32" height="32"> Multi-Step Calculations**: Provides the ability to estimate gradients over multiple simulation steps, by propagating gradients through the entire trajectory.
+This is a finite-difference method. In contact-rich systems the result is a
+**finite-difference estimate through nonsmooth contact dynamics**, so it can be noisy or
+step-size-sensitive at discontinuities.
 
-**<img src="https://cdn-icons-png.flaticon.com/512/12979/12979130.png" width="32" height="32"> Batch Simulation Support**: Enables batched simulations and gradient computations, significantly improving computational efficiency for large-scale experiments.
+## Install
 
+```bash
+pip install .
+```
 
-## Execution Benchmark
-<div style="text-align: center;">
-    <img src="execution_time.svg" alt="Benchmark Results">
-</div>
+The native extension targets Linux and macOS and needs a C++17 compiler plus `ninja`.
 
-## Usage
+## Quickstart
 
 ```python
+import mujoco
 import torch
-import mujoco as mj
-from DiffMjStep import MjStep
 
-# Initialize MuJoCo model and data
-xml_path = 'path/to/your/model.xml'
-mj_model = mj.MjModel.from_xml_path(filename=xml_path)
-mj_data = mj.MjData(mj_model)
+from diffmjstep import mj_step
+
+xml = """
+<mujoco model="pendulum">
+  <option timestep="0.01" integrator="Euler"/>
+  <worldbody>
+    <body name="pole">
+      <joint name="hinge" type="hinge" axis="0 1 0" damping="0.1"/>
+      <geom type="capsule" fromto="0 0 0 0 0 -1" size="0.05" density="1"/>
+    </body>
+  </worldbody>
+  <actuator><motor joint="hinge" gear="1"/></actuator>
+</mujoco>
+"""
+
+model = mujoco.MjModel.from_xml_string(xml)
+x0 = torch.tensor([[0.1, 0.0]], dtype=torch.float64, requires_grad=True)
+u = torch.tensor([[0.2]], dtype=torch.float64, requires_grad=True)
+
+x1 = mj_step(model, x0, u)
+x1.square().sum().backward()
+print(x1, x0.grad, u.grad)
+```
 
-# Define initial state and control input tensors
-state = torch.rand(mj_model.nq + mj_model.nv + mj_model.na, requires_grad=True)
-ctrl = torch.rand(mj_model.nu, requires_grad=True)
+Run the complete example with:
 
-# Compute next state and gradients
-next_state, dydx, dydu = MjStep.apply(state, ctrl, n_steps=4, mj_model, mj_model, mj_data)
+```bash
+python examples/00_basic_step.py
 ```
 
-## Notes
-- As of [MuJoCo 3.1.2](https://mujoco.readthedocs.io/en/3.1.2/changelog.html#python-bindings) the initial state passed to `rollout()` must include a time-step, such that `nstate = mj_stateSize(model, mjtState.mjSTATE_FULLPHYSICS)`. 
+The other examples cover trajectory optimization and sensor losses.
 
-## Citation
+## Rollouts and linearization
+
+`mj_rollout` accepts an initial state `x0` and controls `U` with shapes `[B, nx]` and
+`[B, T, nu]`. With `return_all=True` it returns `[B, T+1, nx]`:
+
+```python
+from diffmjstep import mj_rollout
+
+X = mj_rollout(model, x0, U)
+X[:, -1].square().sum().backward()
+```
 
-If you use this package in your research, a citation would be appreciated:
+For a single transition, including free-joint and quaternion models, use `mj_linearize`:
 
+```python
+from diffmjstep import mj_linearize
+
+A, B = mj_linearize(model, x, u)
+A, B, C, D = mj_linearize(model, x, u, sensors=True)
 ```
- @software{DiffMjStep2024,
+
+`mj_linearize` takes the full `[nq + nv + na]` state and returns `A`/`B` in MuJoCo's
+`[2*nv + na]` tangent space. Autograd rollouts currently require `nq == nv`.
+
+## Sensors
+
+```python
+X, Y = mj_rollout(model, x0, U, return_sensors=True)
+sensor_loss = (Y[:, -1] - target).square().sum()
+sensor_loss.backward()
+```
+
+Sensor timing follows MuJoCo's step contract: `Y[t] = g(x_t, u_t)`. Therefore
+`Y[:, -1]` is the **last pre-step sensor**, not a reading at the terminal state `X[:, -1]`.
+Losses may depend on `X`, `Y`, or both.
+
+## Backends
+
+- `auto` (default): C-threaded forward; native threaded backward when it loads, otherwise
+  the Python MuJoCo transition-FD VJP.
+- `mujoco_rollout`: C-threaded forward and Python backward.
+- `cpp_vjp`: C-threaded forward and required native threaded backward.
+
+The native code is packaged at
+[`diffmjstep/cpp/rollout_vjp.cpp`](diffmjstep/cpp/rollout_vjp.cpp). `nthread` can override
+the conservative automatic thread count.
+
+## Validation and benchmarks
+
+Speedup over a single-threaded pure-Python reference of the *same* algorithm, per dm_control
+model, out of the box (batch 256): forward scales with simplicity (11–23×), backward is a
+steady ~5–6×.
+
+<p align="center">
+  <img src="assets/speedup_per_model.svg" width="820"
+       alt="Per-model DiffMjStep speedup: forward 11-23x, backward 4-6x over a pure-Python reference of the same algorithm (batch 256, MuJoCo 3.10).">
+</p>
+
+Validation and benchmark tooling lives on the
+[`revive/validation` branch](https://github.com/EladSharony/DiffMjStep/tree/revive/validation).
+
+## Scope and limitations
+
+- Autograd state and control inputs must be CPU `float32` or `float64` tensors on the same
+  device.
+- Autograd rollouts require `nq == nv`; `mj_linearize` supports free-joint and quaternion
+  models in tangent space.
+- RK4 is unsupported by `mjd_transitionFD`; use Euler, implicit, or implicitfast.
+- Higher-order gradients are unsupported because the custom backward is
+  `once_differentiable`.
+- DiffMjStep disables MuJoCo solver warm-start on its private runtime model because the
+  warm-start buffer is hidden state outside `x`. This keeps no-grad and autograd forward
+  dynamics identical without mutating the caller's model.
+  Model owners may pre-disable `mjDSBL_WARMSTART` for the zero-copy inference path.
+- Pure callbacks whose outputs depend only on model, state, control, and simulation time are
+  supported. For derivative calls, registrations must not change or be mutated concurrently
+  from rollout forward through backward. An active Python
+  callback uses the Python transition-FD path; behavior that depends on external mutable state
+  is unsupported because finite-difference replay cannot reproduce it reliably.
+- Plugin/history state that extends MuJoCo `FULLPHYSICS` snapshots is unsupported.
+- Contact gradients remain a finite-difference estimate through nonsmooth contact dynamics.
+
+## Citation
+
+```bibtex
+@software{DiffMjStep2026,
   author = {Sharony, Elad},
-  title = {{DiffMjStep: Custom Autograd Function for Differentiable MuJoCo Dynamics}},
-  year = {2024},
-  version = {1.0},
+  title = {{DiffMjStep: PyTorch Autograd for MuJoCo Rollouts}},
+  year = {2026},
+  version = {0.2.0},
   howpublished = {\url{https://github.com/EladSharony/DiffMjStep}},
 }
 ```