diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..0c4c694 --- /dev/null +++ b/.gitignore @@ -0,0 +1,8 @@ +__pycache__/ +*.py[cod] +.venv/ +*.egg-info/ +dist/ +build/ +MUJOCO_LOG.TXT +.idea/ diff --git a/.idea/.gitignore b/.idea/.gitignore deleted file mode 100644 index a9d7db9..0000000 --- a/.idea/.gitignore +++ /dev/null @@ -1,10 +0,0 @@ -# Default ignored files -/shelf/ -/workspace.xml -# Editor-based HTTP Client requests -/httpRequests/ -# Datasource local storage ignored files -/dataSources/ -/dataSources.local.xml -# GitHub Copilot persisted chat sessions -/copilot/chatSessions diff --git a/.idea/DiffMjStep.iml b/.idea/DiffMjStep.iml deleted file mode 100644 index 2368678..0000000 --- a/.idea/DiffMjStep.iml +++ /dev/null @@ -1,10 +0,0 @@ - - - - - - - - - - \ No newline at end of file diff --git a/.idea/inspectionProfiles/Project_Default.xml b/.idea/inspectionProfiles/Project_Default.xml deleted file mode 100644 index 34aac04..0000000 --- a/.idea/inspectionProfiles/Project_Default.xml +++ /dev/null @@ -1,36 +0,0 @@ - - - - \ No newline at end of file diff --git a/.idea/inspectionProfiles/profiles_settings.xml b/.idea/inspectionProfiles/profiles_settings.xml deleted file mode 100644 index 105ce2d..0000000 --- a/.idea/inspectionProfiles/profiles_settings.xml +++ /dev/null @@ -1,6 +0,0 @@ - - - - \ No newline at end of file diff --git a/.idea/misc.xml b/.idea/misc.xml deleted file mode 100644 index 90a19b2..0000000 --- a/.idea/misc.xml +++ /dev/null @@ -1,7 +0,0 @@ - - - - - - \ No newline at end of file diff --git a/.idea/modules.xml b/.idea/modules.xml deleted file mode 100644 index 9f965b1..0000000 --- a/.idea/modules.xml +++ /dev/null @@ -1,8 +0,0 @@ - - - - - - - - \ No newline at end of file diff --git a/.idea/vcs.xml b/.idea/vcs.xml deleted file mode 100644 index 94a25f7..0000000 --- a/.idea/vcs.xml +++ /dev/null @@ -1,6 +0,0 @@ - - - - - - \ No newline at end of file diff --git a/CITATION.cff b/CITATION.cff index e62f44f..de747cc 100644 --- a/CITATION.cff +++ b/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 diff --git a/README.md b/README.md index 0361d4a..53ebe14 100644 --- a/README.md +++ b/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) +

+ 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). +

-## 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 -** 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. -** 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. -** Batch Simulation Support**: Enables batched simulations and gradient computations, significantly improving computational efficiency for large-scale experiments. +## Install +```bash +pip install . +``` -## Execution Benchmark -
- Benchmark Results -
+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 = """ + + +""" + +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×. + +

+ Per-model DiffMjStep speedup: forward 11-23x, backward 4-6x over a pure-Python reference of the same algorithm (batch 256, MuJoCo 3.10). +

+ +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}}, } ``` diff --git a/assets/benchmark.svg b/assets/benchmark.svg new file mode 100644 index 0000000..325ea2b --- /dev/null +++ b/assets/benchmark.svg @@ -0,0 +1,58 @@ + + DiffMjStep: 51x faster forward and 13x faster backward than a pure-Python reference of the same algorithm, at peak batch 4096 (geomean over 9 dm_control envs, float64, MuJoCo 3.10.0, Intel i9-14900K, 2 campaigns). Out of the box at batch 256 it is 16x forward / 5.7x backward. + + + + + + + + DiffMjStep makes differentiable MuJoCo fast + peak throughput at batch 4096 — out of the box (batch 256) it is 16× forward / 5.7× backward + + + + pure-Python reference + + DiffMjStep + + + throughput (transitions / s, log scale) + + + + 10M + + 1M + + 100k + + 10k + + + + 85k + + 4.33M + 51× + + + + + 33k + + 415k + 13× + + + + Forward + Backward + + + throughput = B×T ÷ wall-clock · geomean over 9 dm_control envs (1–9 DOF) · per-env n = 15 fwd / 8 bwd runs · 2 campaigns, spread ≤ 5% + baseline = the same algorithm, no native acceleration — fwd: naive per-env Python mj_step loop; bwd: reference NumPy mjd_transitionFD VJP + Intel i9-14900K · MuJoCo 3.10.0 · PyTorch 2.12.1+cpu · float64 · batch 4096 · horizon 32 · geomean over 9 dm_control envs, 2 campaigns + diff --git a/assets/half_cheetah.xml b/assets/half_cheetah.xml deleted file mode 100644 index 64272d4..0000000 --- a/assets/half_cheetah.xml +++ /dev/null @@ -1,96 +0,0 @@ - - - - - - - - - - \ No newline at end of file diff --git a/assets/speedup_per_model.svg b/assets/speedup_per_model.svg new file mode 100644 index 0000000..98ff48e --- /dev/null +++ b/assets/speedup_per_model.svg @@ -0,0 +1,77 @@ + +DiffMjStep speedup over a pure-Python reference of the same algorithm, per dm_control model (batch 256): forward 11–23x, backward 4–6x. Geomean per model, mujoco 3.10.0, Intel i9-14900K. +DiffMjStep speedup by dm_control model +batch 256 — forward 11–23×, backward 4–6× over a pure-Python reference of the same algorithm + +forward + +backward +speedup over the Python reference + +5× + +10× + +15× + +20× + + +1× = reference + +20× + +5.6× +pendulum +1 dof + +23× + +5.6× +acrobot +2 dof + +20× + +4.4× +cartpole +2 dof + +18× + +5.9× +pointmass +2 dof + +21× + +5.6× +reacher +2 dof + +14× + +5.8× +finger +3 dof + +13× + +5.7× +hopper +7 dof + +12× + +6.4× +cheetah +9 dof + +11× + +6.2× +walker +9 dof +forward = C-threaded mujoco_rollout vs a naive per-env Python mj_step loop · backward = native cpp_vjp vs the reference NumPy mjd_transitionFD VJP · geomean per model +Intel i9-14900K · MuJoCo 3.10.0 · PyTorch 2.12.1+cpu · float64 · batch 256 · horizon 32 · per-env n = 30 forward / 20 backward runs · bars = speedup over the same-algorithm reference + diff --git a/autograd_mujoco.py b/autograd_mujoco.py deleted file mode 100644 index e0902eb..0000000 --- a/autograd_mujoco.py +++ /dev/null @@ -1,178 +0,0 @@ -import numpy as np -import torch -from torch.autograd import Function -from torch.autograd.function import once_differentiable - -import mujoco as mj -from mujoco import rollout - - -class MjStep(Function): - """ - A custom autograd function for the MuJoCo step function. - This is required because the MuJoCo step function is not differentiable. - """ - - @staticmethod - def forward(*args, **kwargs) -> tuple[torch.Tensor, np.ndarray, np.ndarray]: - """ - Forward pass of the MjStep function. - - Args: - *args: Variable length argument list. - **kwargs: Arbitrary keyword arguments. - - Returns: - next_state: The next state after the step. - dydx: The derivative of y with respect to x. - dydu: The derivative of y with respect to u. - - Dimensions key: - B: Batch size - nq: Number of position variables - nv: Number of velocity variables - na: Number of actuator variables - nu: Number of control variables - n_steps: Number of steps in the rollout - - state.shape = [B, nq + nv + na] - ctrl.shape = [B, 1, nu] - - """ - # Extracting the arguments - state, ctrl, n_steps, mj_model, mj_data = args - dydx, dydu = [], [] - device = state.device - compute_grads = state.requires_grad or ctrl.requires_grad - - state, ctrl = state.numpy(force=True), ctrl.numpy(force=True) - - # Repeat the control input for each step in the rollout: [B, 1, nu] -> [B, n_steps, nu] - ctrl = np.repeat(ctrl[:, None, :], n_steps, axis=1) - - if state.shape[-1] == mj.mj_stateSize(mj_model, mj.mjtState.mjSTATE_FULLPHYSICS.value) - 1: - """ - As of MuJoCo 3.1.2 the initial state passed to rollout() must include a time step. - Manually concatenating a time-step to the state vector. - """ - state = np.concatenate([np.zeros_like(state[:, [0]]), state], axis=1) - - # Perform the rollout and get the states - states, _ = mj.rollout.rollout(mj_model, mj_data, state, ctrl) - - # Pop the time-step from state and states (we don't need it for the gradients) - state, states = state[:, 1:], states[:, :, 1:] - - # Reshape the states and control inputs: [B, n_steps, nq + nv + na], [B, n_steps, nu] - states = states.reshape(-1, n_steps, mj_model.nq + mj_model.nv + mj_model.na) - ctrl = ctrl.reshape(-1, n_steps, mj_model.nu) - - if compute_grads: - # Concatenate the initial state with the rest of the states - _states = np.concatenate([state[:, None, :], states[:, :-1, :]], axis=1) - # Set the solver tolerances and disable the warmstart for the solver - mj_model.opt.tolerance = 0 # Disable early termination to make the same number of steps for each FD call - mj_model.opt.ls_tolerance = 1e-18 # Set the line search tolerance to a very low value, for stability - mj_model.opt.disableflags = 2 ** 8 # Disable solver warmstart - for (state_batch, ctrl_batch) in zip(_states, ctrl): - # Initialize the A and B matrices for the approximated linear system - A = np.eye(2 * mj_model.nv + mj_model.na) - B = np.zeros((2 * mj_model.nv + mj_model.na, mj_model.nu)) - - for (_state, _ctrl) in zip(state_batch, ctrl_batch): - # Reset the MuJoCo data and set the state and control - mj.mj_resetData(mj_model, mj_data) - MjStep.set_state_and_ctrl(mj_model, mj_data, _state, _ctrl) - - # Initialize the _A and _B matrices for the approximated linear system - _A = np.zeros((2 * mj_model.nv + mj_model.na, 2 * mj_model.nv + mj_model.na)) - _B = np.zeros((2 * mj_model.nv + mj_model.na, mj_model.nu)) - - # Compute the forward dynamics using MuJoCo's built-in function - mj.mjd_transitionFD(mj_model, mj_data, 1e-8, 1, _A, _B, None, None) - - # Update the A and B matrices - A = np.matmul(_A, A) - B = np.matmul(_A, B) + _B - - # Append the A and B matrices to the lists - dydx.append(A.copy()) - dydu.append(B.copy()) - - # Reset the solver tolerances and enable the warmstart for the solver - mj_model.opt.tolerance = 1e-10 - mj_model.opt.ls_tolerance = 1e-10 - mj_model.opt.disableflags = 0 - - # Get the final state from the rollout - next_state = torch.as_tensor(states[:, -1, :], device=device, dtype=torch.float32) - - # Convert the lists of A and B matrices to numpy arrays - dydx = np.array(dydx) if compute_grads else None - dydu = np.array(dydu) if compute_grads else None - - return next_state, dydx, dydu - - @staticmethod - def setup_context(ctx: Function, inputs: tuple, output: tuple) -> None: - """ - Set up the context for the backward pass. - """ - state, _, _, _, _ = inputs - _, dydx, dydu = output - - dydx = torch.from_numpy(dydx.astype(np.float32)).to(state.device) if dydx is not None else None - dydu = torch.from_numpy(dydu.astype(np.float32)).to(state.device) if dydu is not None else None - - ctx.save_for_backward(dydx, dydu) - - @staticmethod - @once_differentiable - def backward(ctx: Function, grad_output: torch.Tensor, *args) \ - -> tuple[torch.Tensor, torch.Tensor, None, None, None]: - """ - Backward pass of the MjStep function. - - Args: - ctx: The context object where results are saved for backward computation. - grad_output: The output of the forward method. - *args: Variable length argument list. - - Returns: - grad_state: The gradient of the state. - grad_ctrl: The gradient of the control. - None, None, None, None: Placeholder for other gradients that are not computed. - """ - dydx, dydu = ctx.saved_tensors - - if ctx.needs_input_grad[0]: - # dL/dx = dL/dy * dy/dx ([B, 1, nq + nv] * [B, nq + nv, nq + nv] = [B, 1, nq + nv]) - grad_state = torch.bmm(grad_output[:, None, :], dydx).squeeze(1) - else: - grad_state = None - - if ctx.needs_input_grad[1]: - # dL/du = dL/dy * dy/du ([B, 1, nq + nv] * [B, nq + nv, 1] = [B, 1, 1]) - grad_ctrl = torch.bmm(grad_output[:, None, :], dydu).squeeze(1) - - else: - grad_ctrl = None - - return grad_state, grad_ctrl, None, None, None - - @staticmethod - def set_state_and_ctrl(mj_model: mj.MjModel, mj_data: mj.MjData, state: np.ndarray, ctrl: np.ndarray) -> None: - """ - Set the state and control for the MuJoCo model. - """ - mj.mj_resetData(mj_model, mj_data) - np.copyto(mj_data.qpos, state[:mj_model.nq].squeeze()) - np.copyto(mj_data.qvel, state[mj_model.nq:mj_model.nq + mj_model.nv].squeeze()) - np.copyto(mj_data.ctrl, ctrl.squeeze()) - - @staticmethod - def get_state(mj_data: mj.MjData) -> np.ndarray: - """ - Get the state from the MuJoCo data. - """ - return np.concatenate([mj_data.qpos, mj_data.qvel]) diff --git a/diffmjstep/__init__.py b/diffmjstep/__init__.py new file mode 100644 index 0000000..10f1944 --- /dev/null +++ b/diffmjstep/__init__.py @@ -0,0 +1,4 @@ +from .core import mj_rollout, mj_step +from .linearize import mj_linearize + +__all__ = ["mj_linearize", "mj_rollout", "mj_step"] diff --git a/diffmjstep/core.py b/diffmjstep/core.py new file mode 100644 index 0000000..5d088db --- /dev/null +++ b/diffmjstep/core.py @@ -0,0 +1,541 @@ +from __future__ import annotations + +import operator +import os +import warnings +from typing import Any + +import numpy as np +import torch +from torch.autograd.function import once_differentiable + +import mujoco +from mujoco import rollout +from . import native +from .linearize import ( + _transition_fd_at_snapshot, + check_transition_fd_integrator, +) +from .state import ( + infer_state_dims, + snapshot_model, + validate_control, + validate_devices, + validate_eps, + validate_fullphysics, + validate_state, + validate_tensor, +) + + +_SUPPORTED_BACKENDS = {"auto", "mujoco_rollout", "cpp_vjp"} +_AUTO_NATIVE_DISABLED = False +_AUTO_NATIVE_LOADING_PID: int | None = None + + +def _check_backend(backend: str) -> None: + if backend not in _SUPPORTED_BACKENDS: + supported = ", ".join(sorted(_SUPPORTED_BACKENDS)) + raise NotImplementedError(f"backend {backend!r} is unknown; choose from {supported}") + + +def _auto_native_ready() -> bool: + global _AUTO_NATIVE_DISABLED, _AUTO_NATIVE_LOADING_PID + pid = os.getpid() + if _AUTO_NATIVE_DISABLED: + return False + if _AUTO_NATIVE_LOADING_PID == pid: + return False + # No Python call between the checks and assignment: the GIL makes this the single + # loading claim without a lock that could be inherited held by fork. + _AUTO_NATIVE_LOADING_PID = pid + try: + native.load_extension() + except Exception as exc: + _AUTO_NATIVE_DISABLED = True + warnings.warn( + "DiffMjStep could not load the native cpp_vjp backend; " + "backend='auto' will keep MuJoCo's C rollout forward and use the " + "Python VJP. Install Ninja and a C++17 compiler. " + f"Native load failed with {type(exc).__name__}: {exc}", + RuntimeWarning, + stacklevel=2, + ) + return False + else: + return True + finally: + if _AUTO_NATIVE_LOADING_PID == pid: + _AUTO_NATIVE_LOADING_PID = None + +def _python_callbacks() -> tuple[Any | None, ...]: + return ( + mujoco.get_mjcb_act_bias(), + mujoco.get_mjcb_act_dyn(), + mujoco.get_mjcb_act_gain(), + mujoco.get_mjcb_contactfilter(), + mujoco.get_mjcb_control(), + mujoco.get_mjcb_passive(), + mujoco.get_mjcb_sensor(), + mujoco.get_mjcb_time(), + ) + + +def _require_same_python_callbacks(expected: tuple[Any | None, ...]) -> None: + if any( + current is not captured + for current, captured in zip(_python_callbacks(), expected, strict=True) + ): + raise RuntimeError( + "MuJoCo callback registration must remain unchanged from forward through " + "backward; register callbacks before mj_rollout" + ) + + +def _tensor_to_numpy(value: torch.Tensor) -> np.ndarray: + return value.detach().cpu().numpy() + + +def _public_output_from_snapshots( + snapshots: np.ndarray, + nx: int, + return_all: bool, +) -> np.ndarray: + public_states = ( + snapshots[..., 1 : 1 + nx] + if return_all + else snapshots[:, -1, 1 : 1 + nx] + ) + return np.array( + public_states, + dtype=np.float64, + copy=True, + order="C", + ) + + +_AUTO_NTHREAD_CAP = 8 # past this, per-call mjData allocation outweighs added parallelism + + +def _auto_nthread(batch_size: int, nthread: int | None) -> int: + if batch_size == 0: + return 1 + if nthread is not None: + return min(max(1, nthread), batch_size) + if batch_size < 8: + return 1 + return min(_AUTO_NTHREAD_CAP, os.cpu_count() or 1, batch_size) + + +def _rollout_mujoco( + model: Any, + x0: np.ndarray, + U: np.ndarray, + *, + return_all: bool, + sensors: bool, + nthread: int | None, +) -> tuple[np.ndarray, np.ndarray | None, np.ndarray]: + """Fast forward rollout via mujoco.rollout (C-threaded, GIL released).""" + dims = infer_state_dims(model) + nfull = validate_fullphysics(model) + batch_size, horizon, _ = U.shape + init = np.zeros((batch_size, nfull), dtype=np.float64) + init[:, 1:] = x0 # leading element is simulation time = 0 + + snapshots = np.empty((batch_size, horizon + 1, nfull), dtype=np.float64) + snapshots[:, 0, :] = init + if batch_size == 0 or horizon == 0: + X = _public_output_from_snapshots( + snapshots, dims.nx_qpos, return_all + ) + Y = ( + np.empty( + (batch_size, horizon, int(model.nsensordata)), + dtype=np.float64, + ) + if sensors + else None + ) + return X, Y, snapshots + + nthread = _auto_nthread(batch_size, nthread) + if nthread > 1: + data_arg: mujoco.MjData | list[mujoco.MjData] = [ + mujoco.MjData(model) for _ in range(nthread) + ] + else: + data_arg = mujoco.MjData(model) + + states, sensordata = rollout.rollout(model, data_arg, init, U) + snapshots[:, 1:, :] = states + X = _public_output_from_snapshots(snapshots, dims.nx_qpos, return_all) + Y = sensordata if sensors else None + return X, Y, snapshots + + +def _rollout_vjp_numpy( + model: Any, + U: np.ndarray, + full_snapshots: np.ndarray, + grad_X: np.ndarray | None, + grad_Y: np.ndarray | None, + *, + return_all: bool, + needs_control_grad: bool, + eps: float, + centered: bool, +) -> tuple[np.ndarray, np.ndarray | None]: + batch_size, horizon, _ = U.shape + nx = infer_state_dims(model).nx_qpos + grad_x0 = np.zeros((batch_size, nx), dtype=np.float64) + grad_U = np.zeros_like(U) if needs_control_grad else None + if batch_size == 0 or (grad_X is None and grad_Y is None): + return grad_x0, grad_U + if horizon == 0: + if grad_X is not None: + grad_x0[:] = grad_X[:, 0] if return_all else grad_X + return grad_x0, grad_U + + work_data = mujoco.MjData(model) + for batch in range(batch_size): + lambda_x = ( + np.zeros(nx, dtype=np.float64) + if return_all or grad_X is None + else grad_X[batch] + ) + for step in reversed(range(horizon)): + if return_all and grad_X is not None: + lambda_x += grad_X[batch, step + 1] + A, B, C, D = _transition_fd_at_snapshot( + model, + work_data, + full_snapshots[batch, step], + U[batch, step], + sensors=grad_Y is not None, + needs_control_grad=needs_control_grad, + eps=eps, + centered=centered, + ) + gy = None if grad_Y is None else grad_Y[batch, step] + if needs_control_grad: + grad_U[batch, step] = lambda_x @ B + if gy is not None: + grad_U[batch, step] += gy @ D + lambda_x = lambda_x @ A + if gy is not None: + lambda_x += gy @ C + grad_x0[batch] = lambda_x + if return_all and grad_X is not None: + grad_x0[batch] += grad_X[batch, 0] + return grad_x0, grad_U + + +class _MjRolloutFunction(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x0: torch.Tensor, + U: torch.Tensor, + model: Any, + X_np: np.ndarray, + Y_np: np.ndarray | None, + full_snapshots: np.ndarray, + return_all: bool, + eps: float, + centered: bool, + nthread: int | None, + backward_impl: str, + callbacks: tuple[Any | None, ...] | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + ctx.set_materialize_grads(False) + ctx.model = model + ctx.full_snapshots = full_snapshots + ctx.return_all = return_all + ctx.eps = eps + ctx.centered = centered + ctx.nthread = nthread + ctx.backward_impl = backward_impl + ctx.callbacks = callbacks + ctx.save_for_backward(U.detach().cpu()) + + X_t = torch.as_tensor(X_np, dtype=x0.dtype, device=x0.device) + Y_t = ( + torch.as_tensor(Y_np, dtype=x0.dtype, device=x0.device) + if Y_np is not None + else x0.new_empty(0) + ) + return X_t, Y_t + + @staticmethod + @once_differentiable + def backward(ctx, grad_X: torch.Tensor | None, grad_Y: torch.Tensor | None): + (U_cpu,) = ctx.saved_tensors + model = ctx.model + return_all = ctx.return_all + + U_np = U_cpu.numpy().astype(np.float64, copy=False) + full_snapshots = ctx.full_snapshots + callbacks = ctx.callbacks + + device = U_cpu.device + + if grad_Y is not None and not bool(torch.any(grad_Y)): + grad_Y = None + + if grad_X is None and grad_Y is None: + grad_x0_t = U_cpu.new_zeros( + (full_snapshots.shape[0], infer_state_dims(model).nx_qpos) + ) + grad_U_t = ( + torch.zeros_like(U_cpu) + if ctx.needs_input_grad[1] + else None + ) + return ( + grad_x0_t, + grad_U_t, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + if callbacks is None: + raise RuntimeError("internal error: differentiable rollout lost callback state") + _require_same_python_callbacks(callbacks) + + if any(callback is not None for callback in callbacks): + use_native = False + elif ctx.backward_impl == "cpp": + native.load_extension() + use_native = True + elif ctx.backward_impl == "auto": + use_native = _auto_native_ready() + else: + use_native = False + + if use_native: + nthr = _auto_nthread(int(full_snapshots.shape[0]), ctx.nthread) + grad_x0_t, grad_U_t = native.rollout_backward_vjp( + model, + full_snapshots, + U_np, + grad_X, + grad_Y, + return_all=return_all, + needs_control_grad=ctx.needs_input_grad[1], + eps=ctx.eps, + centered=ctx.centered, + nthread=nthr, + ) + grad_x0_t = grad_x0_t.to(dtype=U_cpu.dtype, device=device) + if grad_U_t is not None: + grad_U_t = grad_U_t.to(dtype=U_cpu.dtype, device=device) + else: + grad_X_np = ( + grad_X.detach().cpu().numpy().astype(np.float64, copy=False) + if grad_X is not None + else None + ) + grad_Y_np = ( + grad_Y.detach().cpu().numpy().astype(np.float64, copy=False) + if grad_Y is not None + else None + ) + grad_x0, grad_U = _rollout_vjp_numpy( + model, + U_np, + full_snapshots, + grad_X_np, + grad_Y_np, + return_all=return_all, + needs_control_grad=ctx.needs_input_grad[1], + eps=ctx.eps, + centered=ctx.centered, + ) + + grad_x0_t = torch.as_tensor( + grad_x0, dtype=U_cpu.dtype, device=device + ) + grad_U_t = ( + torch.as_tensor(grad_U, dtype=U_cpu.dtype, device=device) + if grad_U is not None + else None + ) + + _require_same_python_callbacks(callbacks) + + return ( + grad_x0_t, + grad_U_t, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +def mj_rollout( + model: Any, + x0: torch.Tensor, + U: torch.Tensor, + *, + backend: str = "auto", + return_all: bool = True, + return_sensors: bool = False, + eps: float = 1e-8, + centered: bool = True, + nthread: int | None = None, +): + """Roll out MuJoCo dynamics from PyTorch tensors. + + backend="auto" uses the C-threaded mujoco.rollout forward and native threaded + backward when it can be loaded, with a Python VJP fallback. backend="mujoco_rollout" + uses the fast forward with the Python VJP; backend="cpp_vjp" requires native backward. + + With return_sensors=True the call returns (X, Y), where Y is the [B, T, nsensordata] + sensor trajectory, and the loss may depend on X, Y, or both. + """ + validate_tensor("state", x0) + validate_tensor("control", U) + validate_devices(x0, U) + if x0.dtype != U.dtype: + raise TypeError( + "state and control tensors must have the same dtype, got " + f"{x0.dtype} and {U.dtype}" + ) + eps = validate_eps(eps) + _check_backend(backend) + validate_state(model, x0) + validate_control(model, U, allow_sequence=True) + if x0.ndim != 2: + raise ValueError("state must include a batch axis with shape [B, nx]") + if U.ndim != 3: + raise ValueError("control must include batch and time axes with shape [B, T, nu]") + if x0.shape[0] != U.shape[0]: + raise ValueError( + f"state/control batch sizes must match, got {x0.shape[0]} and {U.shape[0]}" + ) + + needs_graph = torch.is_grad_enabled() and ( + x0.requires_grad or U.requires_grad + ) + needs_derivatives = needs_graph + warmstart_disabled = int(mujoco.mjtDisableBit.mjDSBL_WARMSTART) + if needs_derivatives or not (int(model.opt.disableflags) & warmstart_disabled): + # Copy first, then validate the private model so a concurrent caller mutation + # cannot invalidate the model between validation and snapshotting. + runtime_model = snapshot_model(model) + else: + validate_fullphysics(model) + runtime_model = model + + if needs_derivatives or backend == "cpp_vjp": + check_transition_fd_integrator(runtime_model) + + backward_impl = ( + "auto" + if backend == "auto" + else "cpp" + if backend == "cpp_vjp" + else "python" + ) + U_contiguous = U.contiguous() + x0_np = np.asarray(_tensor_to_numpy(x0), dtype=np.float64) + U_np = np.asarray(_tensor_to_numpy(U_contiguous), dtype=np.float64) + + callbacks = _python_callbacks() if needs_derivatives else None + X_np, Y_np, full_snapshots = _rollout_mujoco( + runtime_model, + x0_np, + U_np, + sensors=return_sensors, + return_all=return_all, + nthread=nthread, + ) + if callbacks is not None: + _require_same_python_callbacks(callbacks) + X_out, Y_out = _MjRolloutFunction.apply( + x0, + U_contiguous, + runtime_model, + X_np, + Y_np, + full_snapshots, + return_all, + eps, + centered, + nthread, + backward_impl, + callbacks, + ) + return (X_out, Y_out) if return_sensors else X_out + + +def mj_step( + model: Any, + x0: torch.Tensor, + u: torch.Tensor, + *, + nstep: int = 1, + backend: str = "auto", + return_sensors: bool = False, + eps: float = 1e-8, + centered: bool = True, + nthread: int | None = None, +): + """Apply one action for nstep MuJoCo steps and return the final state. + + With return_sensors=True returns (xT, Y), where Y is the [B, nstep, nsensordata] + sensor trajectory over the repeated steps. + """ + validate_tensor("state", x0) + validate_tensor("control", u) + eps = validate_eps(eps) + + try: + nstep = operator.index(nstep) + except TypeError as exc: + raise TypeError( + f"nstep must be an integer, got {type(nstep).__name__}" + ) from exc + if nstep < 1: + raise ValueError(f"nstep must be >= 1, got {nstep}") + validate_state(model, x0) + validate_control(model, u, allow_sequence=False) + if x0.ndim != 2 or u.ndim != 2: + raise ValueError("state and control must include batch axes [B, nx] and [B, nu]") + if x0.shape[0] != u.shape[0]: + raise ValueError( + f"state/control batch sizes must match, got {x0.shape[0]} and {u.shape[0]}" + ) + + U = u[:, None, :].expand(-1, nstep, -1).contiguous() + + return mj_rollout( + model, + x0, + U, + backend=backend, + return_all=False, + return_sensors=return_sensors, + eps=eps, + centered=centered, + nthread=nthread, + ) + + +__all__ = ["mj_rollout", "mj_step"] diff --git a/diffmjstep/cpp/rollout_vjp.cpp b/diffmjstep/cpp/rollout_vjp.cpp new file mode 100644 index 0000000..e327fe6 --- /dev/null +++ b/diffmjstep/cpp/rollout_vjp.cpp @@ -0,0 +1,269 @@ +// Native reverse-time VJP through a MuJoCo rollout, threaded over the batch. +// +// The forward (and therefore the per-step trajectory) is computed in Python via +// mujoco.rollout; this kernel only does the expensive backward: at each (batch, step) +// it recomputes the transition Jacobians A_t, B_t with mjd_transitionFD and accumulates +// the adjoint. The FD sweep is GIL-bound from Python, so here we release the GIL and +// run one mjData + one set of scratch buffers per worker thread. +// +// Layout (all float64, contiguous): +// full_snapshots : [B, T+1, nfull] full physics state including x0 at t = 0 +// U : [B, T, nu] +// grad_X : optional [B, T+1, nx] when return_all else [B, nx] +// grad_Y : optional [B, T, nsensordata], or an empty sentinel +// returns grad_x0 [B, nx], grad_U [B, T, nu] or an empty sentinel. +// +// Correctness mirrors the Python rollout VJPs; qpos state-mode with nq == nv, so the +// nx = 2*nv+na transition Jacobians act directly on the packed state. + +#include +#include + +#include +#include +#include + +namespace { +void process_chunk(const mjModel* m, mjData* d, int b_begin, int b_end, + const double* full_snapshots, const double* U, const double* grad_X, + const double* grad_Y, double* grad_x0, double* grad_U, + int T, int nx, int nfull, int nu, int ns, + bool return_all, bool needs_control_grad, bool use_state, + bool use_sensors, + double eps, mjtByte centered) { + std::vector A(static_cast(nx) * nx); + std::vector B; + if (needs_control_grad) { + B.resize(static_cast(nx) * nu); + } + std::vector C; + if (use_sensors) { + C.resize(static_cast(ns) * nx); + } + std::vector D; + if (use_sensors && needs_control_grad) { + D.resize(static_cast(ns) * nu); + } + std::vector lam(nx), lam_new(nx); + + for (int b = b_begin; b < b_end; ++b) { + if (return_all || !use_state) { + std::fill(lam.begin(), lam.end(), 0.0); + } else { + const double* g = grad_X + static_cast(b) * nx; + for (int i = 0; i < nx; ++i) lam[i] = g[i]; + } + + for (int t = T - 1; t >= 0; --t) { + if (return_all && use_state) { + const double* g = grad_X + (static_cast(b) * (T + 1) + (t + 1)) * nx; + for (int i = 0; i < nx; ++i) lam[i] += g[i]; + } + + const double* full_snapshot = + full_snapshots + (static_cast(b) * (T + 1) + t) * nfull; + mj_resetData(m, d); + mj_setState(m, d, full_snapshot, mjSTATE_FULLPHYSICS); + const double* u = U + (static_cast(b) * T + t) * nu; + for (int i = 0; i < nu; ++i) d->ctrl[i] = u[i]; + + mj_forward(m, d); + const bool forward_difference_D = + centered && use_sensors && needs_control_grad; + mjd_transitionFD(m, d, eps, centered, A.data(), + needs_control_grad ? B.data() : nullptr, + use_sensors ? C.data() : nullptr, + use_sensors && needs_control_grad && !forward_difference_D + ? D.data() + : nullptr); + if (forward_difference_D) { + // MuJoCo's centered-D ordering workaround: a D-only forward pass also + // preserves one-sided control-limit handling. + mjd_transitionFD(m, d, eps, 0, nullptr, nullptr, nullptr, D.data()); + } + + const double* gy = use_sensors + ? grad_Y + (static_cast(b) * T + t) * ns + : nullptr; + + if (needs_control_grad) { + // grad_U[b, t, k] = sum_i lam[i] * B[i, k] + sum_s gy[s] * D[s, k] + double* gu = grad_U + (static_cast(b) * T + t) * nu; + for (int k = 0; k < nu; ++k) { + double s = 0.0; + for (int i = 0; i < nx; ++i) { + s += lam[i] * B[static_cast(i) * nu + k]; + } + if (use_sensors) { + for (int sensor = 0; sensor < ns; ++sensor) { + s += gy[sensor] * D[static_cast(sensor) * nu + k]; + } + } + gu[k] = s; + } + } + // lam <- lam @ A + gy @ C + for (int j = 0; j < nx; ++j) { + double s = 0.0; + for (int i = 0; i < nx; ++i) s += lam[i] * A[static_cast(i) * nx + j]; + if (use_sensors) { + for (int sensor = 0; sensor < ns; ++sensor) { + s += gy[sensor] * C[static_cast(sensor) * nx + j]; + } + } + lam_new[j] = s; + } + std::swap(lam, lam_new); + } + + double* gx = grad_x0 + static_cast(b) * nx; + if (return_all && use_state) { + const double* g0 = grad_X + static_cast(b) * (T + 1) * nx; + for (int i = 0; i < nx; ++i) gx[i] = lam[i] + g0[i]; + } else { + for (int i = 0; i < nx; ++i) gx[i] = lam[i]; + } + } +} + +} // namespace + +std::tuple rollout_backward_vjp( + torch::Tensor full_snapshots, torch::Tensor U, torch::Tensor grad_X, + torch::Tensor grad_Y, + torch::Tensor data_addresses, int64_t model_addr, int64_t nthread, + bool return_all, bool needs_control_grad, double eps, bool centered) { + TORCH_CHECK(full_snapshots.dtype() == torch::kFloat64, + "full_snapshots must be float64"); + TORCH_CHECK(U.dtype() == torch::kFloat64, "U must be float64"); + TORCH_CHECK(grad_X.dtype() == torch::kFloat64, "grad_X must be float64"); + TORCH_CHECK(grad_Y.dtype() == torch::kFloat64, "grad_Y must be float64"); + TORCH_CHECK(data_addresses.dtype() == torch::kInt64, + "data_addresses must be int64"); + TORCH_CHECK(full_snapshots.device().is_cpu() && U.device().is_cpu() && + grad_X.device().is_cpu() && grad_Y.device().is_cpu() && + data_addresses.device().is_cpu(), + "full_snapshots, U, grad_X, grad_Y, and data_addresses " + "must be CPU tensors"); + TORCH_CHECK(full_snapshots.dim() == 3, + "full_snapshots must have shape [B, T+1, nfull]"); + TORCH_CHECK(U.dim() == 3, "U must have shape [B, T, model.nu]"); + TORCH_CHECK(data_addresses.dim() == 1, + "data_addresses must have shape [nthread]"); + const bool use_state = !(grad_X.dim() == 1 && grad_X.size(0) == 0); + const bool use_sensors = !(grad_Y.dim() == 1 && grad_Y.size(0) == 0); + TORCH_CHECK(!use_state || grad_X.dim() == (return_all ? 3 : 2), + return_all ? "grad_X must have shape [B, T+1, nx]" + : "grad_X must have shape [B, nx]"); + TORCH_CHECK(!use_sensors || grad_Y.dim() == 3, + "grad_Y must have shape [B, T, nsensordata]"); + full_snapshots = full_snapshots.contiguous(); + U = U.contiguous(); + grad_X = grad_X.contiguous(); + grad_Y = grad_Y.contiguous(); + data_addresses = data_addresses.contiguous(); + + const mjModel* m = reinterpret_cast(model_addr); + const int B = static_cast(full_snapshots.size(0)); + const int Tp1 = static_cast(full_snapshots.size(1)); + const int snapshot_width = static_cast(full_snapshots.size(2)); + const int T = static_cast(U.size(1)); + const int nq = m->nq, nv = m->nv, na = m->na, nu = m->nu; + const int nx = 2 * nv + na; + const int ns = m->nsensordata; + const int nfull = mj_stateSize(m, mjSTATE_FULLPHYSICS); + + TORCH_CHECK(nq == nv, "cpp_vjp supports qpos state-mode with nq == nv only"); + TORCH_CHECK(m->opt.integrator != mjINT_RK4, + "mjd_transitionFD does not support RK4 integrator; use Euler or implicit integration"); + TORCH_CHECK(Tp1 == T + 1, "full_snapshots must have T+1 steps along dim 1"); + TORCH_CHECK(snapshot_width == nfull, + "FULLPHYSICS snapshot width mismatch: expected ", nfull, + ", got ", snapshot_width); + TORCH_CHECK(static_cast(U.size(0)) == B && + static_cast(U.size(2)) == nu, + "U must have shape [B, T, model.nu]"); + if (use_state) { + if (return_all) { + TORCH_CHECK(static_cast(grad_X.size(0)) == B && + static_cast(grad_X.size(1)) == Tp1 && + static_cast(grad_X.size(2)) == nx, + "grad_X must have shape [B, T+1, nx]"); + } else { + TORCH_CHECK(static_cast(grad_X.size(0)) == B && + static_cast(grad_X.size(1)) == nx, + "grad_X must have shape [B, nx]"); + } + } + if (use_sensors) { + TORCH_CHECK(static_cast(grad_Y.size(0)) == B && + static_cast(grad_Y.size(1)) == T && + static_cast(grad_Y.size(2)) == ns, + "grad_Y must have shape [B, T, nsensordata]"); + } + + auto opts = torch::TensorOptions().dtype(torch::kFloat64); + torch::Tensor grad_x0 = torch::zeros({B, nx}, opts); + torch::Tensor grad_U = needs_control_grad + ? torch::zeros({B, T, nu}, opts) + : torch::empty({0}, opts); + if (B == 0 || T == 0) { + TORCH_CHECK(data_addresses.size(0) == 0, + "data_addresses must be empty when B or T is zero"); + if (T == 0 && use_state) { + grad_x0.copy_(return_all ? grad_X.select(1, 0) : grad_X); + } + return {grad_x0, grad_U}; + } + + int nthr = std::max(1, nthread); + nthr = std::min(nthr, B); + TORCH_CHECK(static_cast(data_addresses.size(0)) == nthr, + "data_addresses must have shape [nthread]"); + + const double* full_snapshots_p = full_snapshots.data_ptr(); + const double* U_p = U.data_ptr(); + const double* grad_X_p = use_state ? grad_X.data_ptr() : nullptr; + const double* grad_Y_p = use_sensors ? grad_Y.data_ptr() : nullptr; + double* gx0_p = grad_x0.data_ptr(); + double* gU_p = needs_control_grad ? grad_U.data_ptr() : nullptr; + + const mjtByte cflag = centered ? 1 : 0; + + // Python owns and registers these mjData instances for callback compatibility. + // The synchronous extension call keeps their wrappers alive until all workers join. + const int64_t* data_addresses_p = data_addresses.data_ptr(); + std::vector datas; + datas.reserve(nthr); + for (int i = 0; i < nthr; ++i) { + mjData* data = reinterpret_cast(data_addresses_p[i]); + TORCH_CHECK(data != nullptr, "data address must not be null"); + datas.push_back(data); + } + + { + pybind11::gil_scoped_release release; + if (nthr <= 1) { + process_chunk(m, datas[0], 0, B, full_snapshots_p, U_p, grad_X_p, + grad_Y_p, gx0_p, gU_p, T, nx, nfull, nu, ns, return_all, + needs_control_grad, use_state, use_sensors, eps, cflag); + } else { + std::vector threads; + for (int ti = 0; ti < nthr; ++ti) { + const int b0 = ti * B / nthr; + const int b1 = (ti + 1) * B / nthr; + threads.emplace_back(process_chunk, m, datas[ti], b0, b1, + full_snapshots_p, U_p, grad_X_p, grad_Y_p, + gx0_p, gU_p, T, nx, nfull, nu, ns, return_all, + needs_control_grad, use_state, use_sensors, eps, cflag); + } + for (auto& th : threads) th.join(); + } + } + return {grad_x0, grad_U}; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, mod) { + mod.def("rollout_backward_vjp", &rollout_backward_vjp, + "Threaded reverse-time VJP through a MuJoCo rollout (CPU, float64)."); +} diff --git a/diffmjstep/linearize.py b/diffmjstep/linearize.py new file mode 100644 index 0000000..b955847 --- /dev/null +++ b/diffmjstep/linearize.py @@ -0,0 +1,191 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np + +import mujoco +from .state import ( + FULLPHYSICS, + as_numpy, + infer_state_dims, + set_data_control, + set_data_state, + snapshot_model, + validate_control, + validate_eps, + validate_state, + validate_tensor, +) + + +def _enum_value(value: Any) -> int: + return int(value.value) if hasattr(value, "value") else int(value) + + +def check_transition_fd_integrator(model: Any) -> None: + if _enum_value(model.opt.integrator) == _enum_value(mujoco.mjtIntegrator.mjINT_RK4): + raise ValueError("mjd_transitionFD does not support RK4 integrator; use Euler or implicit integration") + + +def _as_real_float64(name: str, value: Any) -> np.ndarray: + if hasattr(value, "detach"): + validate_tensor(name, value) + array = as_numpy(value) + else: + array = np.asarray(value) + if isinstance(value, np.ndarray) and not np.issubdtype(array.dtype, np.floating): + raise TypeError(f"{name} must be a real floating array, got {array.dtype}") + if not isinstance(value, np.ndarray) and array.dtype.kind not in {"i", "u", "f"}: + raise TypeError(f"{name} must contain real numeric values, got {array.dtype}") + return np.asarray(array, dtype=np.float64) + + +def _as_batched_state(model: Any, x: Any) -> np.ndarray: + x_np = _as_real_float64("state", x) + validate_state(model, x_np, require_square=False) + return x_np[None, :] if x_np.ndim == 1 else x_np + + +def _as_batched_control(model: Any, u: Any) -> np.ndarray: + u_np = _as_real_float64("control", u) + validate_control(model, u_np, allow_sequence=False) + return u_np[None, :] if u_np.ndim == 1 else u_np + + +def _return_like_input(reference: Any, value: np.ndarray): + if hasattr(reference, "detach"): + import torch + + return torch.as_tensor(value, dtype=reference.dtype, device=reference.device) + return value + + +def _transition_fd_at_snapshot( + model: Any, + data: Any, + full_snapshot: np.ndarray, + u: np.ndarray, + *, + sensors: bool = False, + needs_control_grad: bool = True, + eps: float, + centered: bool, +): + nx = 2 * int(model.nv) + int(model.na) + nu = int(model.nu) + ns = int(model.nsensordata) + A = np.empty((nx, nx), dtype=np.float64) + B = np.empty((nx, nu), dtype=np.float64) if needs_control_grad else None + C = np.empty((ns, nx), dtype=np.float64) if sensors else None + D = ( + np.empty((ns, nu), dtype=np.float64) + if sensors and needs_control_grad + else None + ) + + mujoco.mj_resetData(model, data) + mujoco.mj_setState(model, data, full_snapshot, FULLPHYSICS) + data.ctrl[:] = u + mujoco.mj_forward(model, data) + forward_difference_D = centered and D is not None + mujoco.mjd_transitionFD( + model, + data, + eps, + int(centered), + A, + B, + C, + None if forward_difference_D else D, + ) + if forward_difference_D: + # MuJoCo's centered-D ordering reverses interior control derivatives; a + # D-only forward pass also preserves its one-sided control-limit handling. + mujoco.mjd_transitionFD(model, data, eps, 0, None, None, None, D) + return A, B, C, D + + +def _linearize( + model: Any, + x: Any, + u: Any, + *, + sensors: bool = False, + eps: float = 1e-8, + centered: bool = True, +): + dims = infer_state_dims(model) + x_np = _as_batched_state(model, x) + nx = dims.nx_tangent + + u_np = _as_batched_control(model, u) + if x_np.shape[0] != u_np.shape[0]: + raise ValueError(f"state/control batch sizes must match, got {x_np.shape[0]} and {u_np.shape[0]}") + batch_size = x_np.shape[0] + A = np.empty((batch_size, nx, nx), dtype=np.float64) + B = np.empty((batch_size, nx, dims.nu), dtype=np.float64) + C = np.empty((batch_size, int(model.nsensordata), nx), dtype=np.float64) if sensors else None + D = np.empty((batch_size, int(model.nsensordata), dims.nu), dtype=np.float64) if sensors else None + + data = mujoco.MjData(model) + centered_flag = 1 if centered else 0 + for batch in range(batch_size): + mujoco.mj_resetData(model, data) + set_data_state(model, data, x_np[batch]) + set_data_control(model, data, u_np[batch]) + mujoco.mj_forward(model, data) + forward_difference_D = sensors and centered + mujoco.mjd_transitionFD( + model, + data, + eps, + centered_flag, + A[batch], + B[batch], + C[batch] if sensors else None, + None if forward_difference_D else D[batch] if sensors else None, + ) + if forward_difference_D: + # MuJoCo's centered-D ordering workaround: request D alone with forward + # differences, which also keeps one-sided control-limit behavior intact. + mujoco.mjd_transitionFD( + model, data, eps, 0, None, None, None, D[batch] + ) + + A_out = _return_like_input(x, A) + B_out = _return_like_input(x, B) + if sensors: + return A_out, B_out, _return_like_input(x, C), _return_like_input(x, D) + return A_out, B_out + + +def mj_linearize( + model: Any, + x: Any, + u: Any, + *, + sensors: bool = False, + eps: float = 1e-8, + centered: bool = True, +): + """Compute MuJoCo transition Jacobians A/B, and optionally sensor Jacobians C/D. + + Returned tensors/arrays are batched: A has shape [B, nx, nx], B has shape [B, nx, nu]. + + For free-joint / quaternion models (nq != nv), the input state is the full + [nq + nv + na] vector and A/B are returned in MuJoCo's tangent space (nx = 2*nv + na) -- + the standard linearization for trajopt/iLQR. The qpos-square form is only defined when + nq == nv. + """ + private_model = snapshot_model(model) + eps = validate_eps(eps) + check_transition_fd_integrator(private_model) + return _linearize( + private_model, + x, + u, + sensors=sensors, + eps=eps, + centered=centered, + ) diff --git a/diffmjstep/native.py b/diffmjstep/native.py new file mode 100644 index 0000000..65e2ca2 --- /dev/null +++ b/diffmjstep/native.py @@ -0,0 +1,142 @@ +"""Lazy JIT build + thin wrapper for the native rollout VJP backend. + +The C++ extension is compiled on first use via torch.utils.cpp_extension.load and +cached under TORCH_EXTENSIONS_DIR; subsequent imports are instant. It links the +MuJoCo shared library that ships with the `mujoco` wheel, so no system MuJoCo install +is required. Building needs a C++17 compiler. +""" +from __future__ import annotations + +import functools +import os +import shlex +import sys +from importlib.resources import as_file, files +from pathlib import Path +from typing import Any + +import torch + + +def _ensure_ninja_on_path() -> None: + # torch's JIT build shells out to the `ninja` binary; pip installs it next to the + # interpreter, which is not on PATH unless the venv is activated. Add it so the + # extension builds whether or not the environment was sourced. + bin_dir = os.path.dirname(sys.executable) + parts = os.environ.get("PATH", "").split(os.pathsep) + if bin_dir and bin_dir not in parts: + os.environ["PATH"] = bin_dir + os.pathsep + os.environ.get("PATH", "") + + +def _find_mujoco_library(mujoco_dir: Path, version: str) -> Path: + if sys.platform == "linux": + name = f"libmujoco.so.{version}" + elif sys.platform == "darwin": + name = f"libmujoco.{version}.dylib" + else: + raise RuntimeError( + f"native cpp_vjp does not support {sys.platform!r}; " + "supported targets are Linux and macOS" + ) + library = mujoco_dir / name + if not library.is_file(): + raise RuntimeError(f"could not find MuJoCo shared library {library}") + return library + + +@functools.lru_cache(maxsize=1) +def load_extension(): + import mujoco + from torch.utils.cpp_extension import load + + _ensure_ninja_on_path() + mujoco_dir = Path(mujoco.__file__).resolve().parent + source_resource = files("diffmjstep") / "cpp" / "rollout_vjp.cpp" + library = _find_mujoco_library(mujoco_dir, version=mujoco.__version__) + include = mujoco_dir / "include" + if not include.is_dir(): + raise RuntimeError(f"MuJoCo headers are missing: {include}") + + pthread = ["-pthread"] if sys.platform == "linux" else [] + with as_file(source_resource) as source: + if not source.is_file(): + raise RuntimeError(f"packaged native source is missing: {source}") + return load( + name="diffmjstep_cpp", + sources=[str(source)], + extra_include_paths=[str(include)], + extra_cflags=["-O3", "-std=c++17", *pthread], + extra_ldflags=[ + shlex.quote(str(library)), + shlex.quote(f"-Wl,-rpath,{mujoco_dir}"), + *pthread, + ], + verbose=False, + ) + + +def rollout_backward_vjp( + model: Any, + full_snapshots: Any, + U: Any, + grad_X: Any | None, + grad_Y: Any | None = None, + *, + return_all: bool, + needs_control_grad: bool, + eps: float, + centered: bool, + nthread: int, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Native reverse-time VJP. + + full_snapshots: [B, T+1, nfull], U: [B, T, nu]. grad_X is optional + [B, T+1, nx] or [B, nx], and grad_Y is optional [B, T, nsensordata]. Returns + (grad_x0 [B, nx], grad_U [B, T, nu] or None). + """ + import mujoco + + ext = load_extension() + full_snapshots = ( + torch.as_tensor(full_snapshots).detach().to(torch.float64).cpu().contiguous() + ) + U = torch.as_tensor(U).detach().to(torch.float64).cpu().contiguous() + grad_X = ( + torch.empty(0, dtype=torch.float64, device="cpu") + if grad_X is None + else torch.as_tensor(grad_X).detach().to(torch.float64).cpu().contiguous() + ) + grad_Y = ( + torch.empty(0, dtype=torch.float64, device="cpu") + if grad_Y is None + else torch.as_tensor(grad_Y).detach().to(torch.float64).cpu().contiguous() + ) + has_work = ( + full_snapshots.ndim == 3 + and U.ndim == 3 + and full_snapshots.shape[0] > 0 + and U.shape[1] > 0 + ) + worker_count = ( + min(max(1, int(nthread)), int(full_snapshots.shape[0])) + if has_work + else 0 + ) + scratch_data = [mujoco.MjData(model) for _ in range(worker_count)] + data_addresses = torch.tensor( + [data._address for data in scratch_data], dtype=torch.int64, device="cpu" + ) + grad_x0, grad_U = ext.rollout_backward_vjp( + full_snapshots, + U, + grad_X, + grad_Y, + data_addresses, + int(model._address), + int(nthread), + bool(return_all), + bool(needs_control_grad), + float(eps), + bool(centered), + ) + return grad_x0, grad_U if needs_control_grad else None diff --git a/diffmjstep/state.py b/diffmjstep/state.py new file mode 100644 index 0000000..e2e120c --- /dev/null +++ b/diffmjstep/state.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +import copy +import math +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch + +import mujoco + + +FULLPHYSICS = mujoco.mjtState.mjSTATE_FULLPHYSICS + + +def validate_tensor(name: str, value: object) -> None: + if not isinstance(value, torch.Tensor): + raise TypeError( + f"{name} must be a float32 or float64 torch.Tensor, " + f"got {type(value).__name__}" + ) + if value.layout != torch.strided: + raise TypeError(f"{name} must have strided layout, got {value.layout}") + if value.dtype not in {torch.float32, torch.float64}: + raise TypeError(f"{name} must be float32 or float64, got {value.dtype}") + + +def validate_devices(state: torch.Tensor, control: torch.Tensor) -> None: + if state.device != control.device: + raise ValueError( + "state and control tensors must be on the same device, got " + f"{state.device} and {control.device}" + ) + if state.device.type != "cpu": + raise ValueError(f"state and control must be on CPU, got {state.device}") + + +def validate_eps(eps: float) -> float: + eps = float(eps) + if not math.isfinite(eps) or eps <= 0: + raise ValueError(f"eps must be finite and positive, got {eps!r}") + return eps + + +def validate_fullphysics(model: mujoco.MjModel) -> int: + actual = int(mujoco.mj_stateSize(model, FULLPHYSICS)) + expected = 1 + int(model.nq) + int(model.nv) + int(model.na) + if actual != expected: + raise NotImplementedError( + "plugin/history FULLPHYSICS state is unsupported: " + f"expected {expected} values (time+qpos+qvel+act), got {actual}" + ) + return actual + + +def snapshot_model(model: mujoco.MjModel) -> mujoco.MjModel: + private = copy.copy(model) + validate_fullphysics(private) + private.opt.disableflags |= int(mujoco.mjtDisableBit.mjDSBL_WARMSTART) + return private + + +@dataclass(frozen=True) +class StateDims: + nq: int + nv: int + na: int + nu: int + nx_qpos: int + nx_tangent: int + + +def infer_state_dims(model: Any) -> StateDims: + """Return the state dimensions DiffMjStep uses for the given MuJoCo model.""" + nq = int(model.nq) + nv = int(model.nv) + na = int(model.na) + nu = int(model.nu) + return StateDims( + nq=nq, + nv=nv, + na=na, + nu=nu, + nx_qpos=nq + nv + na, + nx_tangent=2 * nv + na, + ) + + +def _shape(value: Any) -> tuple[int, ...]: + if not hasattr(value, "shape"): + raise TypeError("expected an array-like object with a shape") + return tuple(int(dim) for dim in value.shape) + + +def check_rollout_state_supported(model: Any) -> None: + dims = infer_state_dims(model) + if dims.nq != dims.nv: + raise NotImplementedError( + "rollout state requires nq == nv; free-joint and quaternion models " + "are supported by mj_linearize only" + ) + + +def validate_state(model: Any, x: Any, *, require_square: bool = True) -> None: + """Validate a state tensor/array shape; raise ValueError if it is wrong.""" + shape = _shape(x) + if len(shape) not in (1, 2): + raise ValueError(f"state must have shape [nx] or [B, nx], got {shape}") + + if require_square: + check_rollout_state_supported(model) + nx = infer_state_dims(model).nx_qpos + if shape[-1] != nx: + raise ValueError(f"state last dimension must be {nx}, got {shape[-1]}") + + +def validate_control(model: Any, u: Any, allow_sequence: bool = True) -> None: + """Validate a control tensor/array shape; raise ValueError if it is wrong.""" + shape = _shape(u) + allowed = (1, 2, 3) if allow_sequence else (1, 2) + if len(shape) not in allowed: + expected = "[nu], [B, nu], or [B, T, nu]" if allow_sequence else "[nu] or [B, nu]" + raise ValueError(f"control must have shape {expected}, got {shape}") + + nu = int(model.nu) + if shape[-1] != nu: + raise ValueError(f"control last dimension must be {nu}, got {shape[-1]}") + + +def as_numpy(value: Any) -> np.ndarray: + """Convert a CPU tensor/array-like object to a detached NumPy array.""" + if hasattr(value, "detach"): + if getattr(value, "device", None) is not None and value.device.type != "cpu": + raise ValueError("MuJoCo CPU backend received a non-CPU tensor") + return value.detach().cpu().numpy() + return np.asarray(value) + + +def set_data_state(model: Any, data: Any, x: np.ndarray) -> None: + """Set qpos/qvel/act in mjData from a full [nq+nv+na] DiffMjStep state vector. + + Slicing qpos/qvel/act is valid for any model, including free joints. Rollout's + nq == nv restriction is enforced by validate_state at the public boundary. + """ + if x.ndim != 1: + raise ValueError(f"state must have shape [nx], got {_shape(x)}") + validate_state(model, x, require_square=False) + x = np.asarray(x, dtype=np.float64) + dims = infer_state_dims(model) + + data.qpos[:] = x[: dims.nq] + data.qvel[:] = x[dims.nq : dims.nq + dims.nv] + if dims.na: + data.act[:] = x[dims.nq + dims.nv : dims.nq + dims.nv + dims.na] + + +def set_data_control(model: Any, data: Any, u: np.ndarray) -> None: + u = np.asarray(u, dtype=np.float64) + if u.shape[-1] != int(model.nu): + raise ValueError(f"control last dimension must be {model.nu}, got {u.shape[-1]}") + data.ctrl[:] = u diff --git a/examples/00_basic_step.py b/examples/00_basic_step.py new file mode 100644 index 0000000..f63fa16 --- /dev/null +++ b/examples/00_basic_step.py @@ -0,0 +1,39 @@ +import torch +import mujoco + +from diffmjstep import mj_step + + +PENDULUM_XML = """ + + +""" + + +def main() -> None: + model = mujoco.MjModel.from_xml_string(PENDULUM_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, nstep=1) + loss = x1.square().sum() + loss.backward() + + print("x1:", x1.detach().numpy()) + print("d loss / d x0:", x0.grad.numpy()) + print("d loss / d u:", u.grad.numpy()) + + +if __name__ == "__main__": + main() diff --git a/examples/03_trajectory_optimization.py b/examples/03_trajectory_optimization.py new file mode 100644 index 0000000..d70983f --- /dev/null +++ b/examples/03_trajectory_optimization.py @@ -0,0 +1,50 @@ +import torch +import mujoco + +from diffmjstep import mj_rollout + + +PENDULUM_XML = """ + + +""" + + +def main() -> None: + torch.manual_seed(0) + + model = mujoco.MjModel.from_xml_string(PENDULUM_XML) + + horizon = 12 + x0 = torch.tensor([[0.1, 0.0]], dtype=torch.float64) + target = torch.tensor([[0.5, 0.0]], dtype=torch.float64) + U = torch.zeros(1, horizon, model.nu, dtype=torch.float64, requires_grad=True) + optimizer = torch.optim.Adam([U], lr=0.05) + + for _ in range(25): + optimizer.zero_grad() + X = mj_rollout(model, x0, U, return_all=True) + terminal_loss = (X[:, -1] - target).square().sum() + control_loss = 1e-3 * U.square().sum() + loss = terminal_loss + control_loss + loss.backward() + optimizer.step() + + X = mj_rollout(model, x0, U, return_all=True) + print("final state:", X[:, -1].detach().numpy()) + print("target:", target.numpy()) + print("optimized controls:", U.detach().numpy()) + + +if __name__ == "__main__": + main() diff --git a/examples/04_sensor_loss.py b/examples/04_sensor_loss.py new file mode 100644 index 0000000..1e5cdcb --- /dev/null +++ b/examples/04_sensor_loss.py @@ -0,0 +1,58 @@ +"""Differentiate a loss on the sensor trajectory: drive the pendulum tip to a target. + +The framepos sensor gives the tip's world position; we optimize controls so the tip +reaches a target, using gradients that flow through MuJoCo's sensor Jacobians (C/D). +""" +import torch +import mujoco + +from diffmjstep import mj_rollout + +SENSORS_XML = """ + + +""" + + +def main() -> None: + model = mujoco.MjModel.from_xml_string(SENSORS_XML) + + # sensordata layout: [jointpos(1), jointvel(1), framepos(3)]; tip xyz is the last 3. + horizon = 12 + x0 = torch.tensor([[0.1, 0.0]], dtype=torch.float64) + # reachable point on the radius-1 tip circle: (sin 0.6, 0, -cos 0.6) + target_xyz = torch.tensor([0.5646, 0.0, -0.8253], dtype=torch.float64) + U = torch.zeros(1, horizon, model.nu, dtype=torch.float64, requires_grad=True) + opt = torch.optim.Adam([U], lr=0.05) + + for _ in range(40): + opt.zero_grad() + _, Y = mj_rollout(model, x0, U, return_sensors=True, return_all=True) + tip = Y[:, -1, -3:] + loss = (tip - target_xyz).square().sum() + 1e-3 * U.square().sum() + loss.backward() + opt.step() + + _, Y = mj_rollout(model, x0, U, return_sensors=True, return_all=True) + print("last pre-step tip xyz:", Y[0, -1, -3:].detach().numpy()) + print("target tip xyz:", target_xyz.numpy()) + + +if __name__ == "__main__": + main() diff --git a/execution_time.svg b/execution_time.svg deleted file mode 100644 index 93e1898..0000000 --- a/execution_time.svg +++ /dev/null @@ -1,1469 +0,0 @@ - - - - - - - - 2024-03-30T19:17:47.754564 - image/svg+xml - - - Matplotlib v3.8.0, https://matplotlib.org/ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/main.py b/main.py deleted file mode 100644 index e8fa6d4..0000000 --- a/main.py +++ /dev/null @@ -1,104 +0,0 @@ -import torch -import mujoco as mj -from autograd_mujoco import MjStep - -import time -import pandas as pd -import matplotlib.pyplot as plt -import seaborn as sns - - -def jacobian(f: callable, x: torch.Tensor, eps: float = 1e-4) -> torch.Tensor: - """ - Computes the Jacobian of a function f at x using centered finite differences. - :param f: function - :param x: input - :param eps: epsilon - :return: Jacobian of f at x - """ - x = x.view(1, -1) - n = x.numel() - e = torch.eye(n, dtype=torch.float64, device=x.device) - J = [(f(x + eps * e[i]) - f(x - eps * e[i])) / (2 * eps) for i in range(n)] - return torch.stack(J).reshape(n, -1).T - - -def call_jacobian(state: torch.Tensor, ctrl: torch.Tensor, n_steps: int, mj_model: mj.MjModel, mj_data: mj.MjData) -> \ - tuple[torch.Tensor, torch.Tensor]: - """ - Uses the jacobian function to compute the dydx and dydu. - """ - state.requires_grad = False - ctrl.requires_grad = False - n_batch = state.shape[0] - naive_dydx, naive_dydu = zip( - *[(jacobian(lambda s: MjStep.apply(s, ctrl[[i]], n_steps, mj_model, mj_data)[0], state[[i]], 1e-4), - jacobian(lambda a: MjStep.apply(state[[i]], a, n_steps, mj_model, mj_data)[0], ctrl[[i]], 1e-4)) - for i in range(n_batch)]) - naive_dydx, naive_dydu = torch.stack(naive_dydx), torch.stack(naive_dydu) - return naive_dydx, naive_dydu - - -if __name__ == '__main__': - # path to the xml file - xml_path = 'assets/half_cheetah.xml' - - # Create an instance of the MjModel and MjData - mj_model = mj.MjModel.from_xml_path(filename=xml_path) - mj_data = mj.MjData(mj_model) - - n_steps = 4 # Number of steps to unroll the dynamics - - torch.manual_seed(0) - - batch_sizes = [2 ** i for i in range(0, 12 + 1)] - - devices = [torch.device('cpu'), torch.device('cuda')] - - df = pd.DataFrame(columns=['device', 'method', 'batch_size', 'execution_time']) - - # Define the number of runs - n_runs = 5 - for run in range(n_runs): - for device in devices: - for n_batch in batch_sizes: - # Define the state and control - state = torch.rand(n_batch, mj_model.nq + mj_model.nv + mj_model.na, device=device, requires_grad=True) - ctrl = torch.rand(n_batch, mj_model.nu, device=device, requires_grad=True) - - # Measure the execution time of MjStep.apply - start_time = time.time() - next_state, dydx, dydu = MjStep.apply(state, ctrl, n_steps, mj_model, mj_data) - end_time = time.time() - new_row = pd.DataFrame({'run': [run], 'device': [device], 'method': ['MjStep'], 'batch_size': [n_batch], - 'execution_time': [end_time - start_time]}) - df = pd.concat([df, new_row], ignore_index=True) - - # Measure the execution time of naive FD - start_time = time.time() - naive_dydx, naive_dydu = call_jacobian(state, ctrl, n_steps, mj_model, mj_data) - end_time = time.time() - new_row = pd.DataFrame( - {'run': [run], 'device': [device], 'method': ['Naive FD'], 'batch_size': [n_batch], - 'execution_time': [end_time - start_time]}) - df = pd.concat([df, new_row], ignore_index=True) - - # Plot the execution times with error bars - sns.set_style('whitegrid') - sns.set_context('paper', font_scale=1.2) - sns.set_palette('deep') - - ax = sns.lineplot(data=df, x='batch_size', y='execution_time', hue='method', style='device', errorbar='sd') - - sns.move_legend(ax, "upper center", ncol=2) - ax.set_xticks(batch_sizes, [f'{batch_size}' for batch_size in batch_sizes]) - ax.set_xlim([batch_sizes[0], batch_sizes[-1]]) - ax.set_xscale('log', base=2) - ax.set_title(f'Execution Time of MjStep and Naive FD (Over {n_runs} Runs)') - ax.set_xlabel('Batch Size') - ax.set_ylabel('Execution Time (s)') - ax.grid(True) - - # Save the plot to svg - plt.savefig('execution_time.svg', format='svg') - plt.savefig('execution_time.png', format='png') diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..9b33b21 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,38 @@ +[build-system] +requires = ["setuptools>=77"] +build-backend = "setuptools.build_meta" + +[project] +name = "diffmjstep" +version = "0.2.0" +description = "Minimal PyTorch autograd helpers for differentiable MuJoCo stepping." +readme = "README.md" +requires-python = ">=3.10" +authors = [{ name = "Elad Sharony" }] +license = "MIT" +license-files = ["LICENSE"] +keywords = ["mujoco", "pytorch", "differentiable-simulation", "optimal-control"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Operating System :: MacOS", + "Operating System :: POSIX :: Linux", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.12", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] +dependencies = [ + "mujoco>=3.5.0", + "ninja>=1.11", + "numpy>=1.24", + "torch>=2.4", +] + +[project.urls] +Repository = "https://github.com/EladSharony/DiffMjStep" +Issues = "https://github.com/EladSharony/DiffMjStep/issues" + +[tool.setuptools.packages.find] +include = ["diffmjstep*"] + +[tool.setuptools.package-data] +diffmjstep = ["cpp/*.cpp"]