Core concept

MuJoCo can be used from outside the step in four ways:

This lesson builds an LQR controller from MuJoCo’s finite-difference derivatives, shows what a control callback does to those derivatives, runs a first-party plugin, and compiles a C program whose result matches Python’s to the last bit.

[!established] Finite-difference derivatives in 3.14.0 mjd_transitionFD(m, d, eps, flg_centered, A, B, C, D) returns the discrete-time linearization of one mj_step around the current state: $\delta x_{t+1} = A\,\delta x_t + B\,\delta u_t$, with the state in the tangent space (positions as mj_differentiatePos differences, then velocities and activations), and optionally the sensor Jacobians $C$, $D$. It differentiates by stepping perturbed copies. mjd_stepFD (src/engine/engine_derivative_fd.c, line 300) perturbs controls and activations with mj_stepSkip(..., mjSTAGE_VEL, ...), velocities with mjSTAGE_POS, and positions with no skip. That is Lesson 20.1’s mj_forwardSkip at work.

Visual intuition

There is no lab for this lesson: everything here is Python and C on your machine. The cart-pole planner demo is the derivative-free alternative to this lesson’s LQR: it samples futures instead of linearizing one.

Mathematics

A linearization you can check

For the cart-pole balanced upright (state: cart position, pole angle, cart velocity, pole rate; timestep 2 ms), centred differences with $\epsilon = 10^{-6}$ give

\[A = \begin{bmatrix} 1.000000 & -0.000005 & 0.002000 & 0.000000\\ 0 & 1.000110 & 0.000001 & 0.002000\\ 0 & -0.002747 & 0.999810 & 0.000009\\ 0 & 0.054946 & 0.000467 & 0.999813\end{bmatrix},\qquad B = \begin{bmatrix} 3.80\times 10^{-6}\\ -9.33\times 10^{-6}\\ 1.90\times 10^{-3}\\ -4.67\times 10^{-3}\end{bmatrix}.\]

The diagonal is 1 plus the step’s effect, the 0.002 entries are the timestep moving velocity into position, and the 0.0549 entry is gravity tipping the pole over. Together with the 0.002 that turns pole rate into angle, it gives $A$ an eigenvalue of about $1 + \sqrt{0.002 \times 0.0549} \approx 1.0105$ per step, above 1: the instability. To check the matrix, step once from a perturbation $\delta x$ of size $s$ and compare the change with $A\,\delta x$. The relative error is $1.2\times 10^{-6}$ for $s = 0.1$, $1.2\times 10^{-8}$ for $s = 0.01$, and falls by a factor of 100 per decade down to rounding ($3\times 10^{-15}$ from $s = 10^{-5}$). It shrinks as $s^2$: the linear model is right and the first missing term is cubic, as expected for a pendulum around a symmetric point. At that symmetric point, forward and centred differences give the same $A$ exactly. At a moving state (0.2, 0.6, 0.5, −1.0) they differ by $2.1\times 10^{-8}$ with $\epsilon = 10^{-6}$ and by $2.1\times 10^{-6}$ with $\epsilon = 10^{-4}$: the first-order error of forward differences.

LQR from MuJoCo’s linearization

With $Q = \text{diag}(1, 10, 1, 1)$ and $R = 0.1$, the discrete Riccati equation gives $K$ = [−3.11, −45.58, −5.68, −9.18] and a closed loop whose largest eigenvalue has magnitude 0.99734, inside the unit circle. On the nonlinear cart-pole, with the actuator’s ±20 N limit, $u = -Kx$ balances the pole from 0.1, 0.3, 0.5 and 0.7 rad, commanding at most 4.6, 13.7, 22.8 and 31.9 N; from 0.9 rad it fails. A linear controller from one linearization has a finite region where it works, and the saturating actuator shrinks it.

Callbacks are global, and the derivatives see them

mjcb_control (in Python, mujoco.set_mjcb_control) is called inside mj_forward before actuation (Lesson 20.1). Two properties follow. First, a Python callback costs Python time on every step: 10.05 µs per cart-pole step against 2.70 µs without. Second, it is a process-wide function pointer, called for every model and every mjData, including the perturbed copies inside mjd_transitionFD. With the LQR callback installed, mjd_transitionFD called it 11 times: once for the nominal step and twice per control, velocity and position perturbation. The control perturbation was overwritten each time, so $B$ came back as exactly zero. $A$ came back as the closed-loop $A - BK$, to $1.4\times 10^{-13}$. Nothing warns. The same holds for mujoco.rollout threads and for any other model in the process. Install callbacks only around the code that needs them, and remove them with set_mjcb_control(None).

Plugins

A plugin is compiled code registered under a name, with its own state, which MJCF attaches to an actuator, sensor, passive force or geometry. The Python package ships four plugin libraries:

Library Plugins
libactuator.so mujoco.pid
libelasticity.so mujoco.elasticity.cable
libsdf_plugin.so mujoco.sdf.bolt, mujoco.sdf.bowl, mujoco.sdf.gear, mujoco.sdf.nut, mujoco.sdf.torus
libsensor.so mujoco.sensor.touch_grid

mujoco.pid with $k_p$ = 40, $k_i$ = 20, $k_d$ = 4 holds a 1 kg pendulum at the commanded 0.5 rad. After 6 s it is at 0.4997 rad, with the integral term’s state (0.0583, in act) carrying the gravity load that a PD controller would leave as an offset (Lesson 8.1). Unlike a callback, a plugin is part of the model: it is declared in MJCF and saved with it, and its state lives in mjData (here in act), so each mjData, and each thread using its own, has its own copy. Writing one means C or C++ against mjpPlugin and building a shared library.

The C API

c/minimal_step.c loads a model, resets it to its first keyframe, steps it 1000 times and prints the positions. Compiled with the headers and the libmujoco that ship inside the Python package, it printed the cart-pole’s qpos as −0.050354155334510445 and 1.899214965103317. Python, stepping the same model the same way, printed the same digits: a difference of 0.0. The Python bindings are a thin layer over the same library.

[!recommendation] Which extension point Use mjd_transitionFD for linearizations, and check them as part (1) does before building a controller on them. Use callbacks only for experiments in one process, and remove them before computing derivatives or running threads. Use a plugin for anything that must travel with the model, keep state, or run in parallel. Drop to C when the Python loop itself is the cost; for most workloads, mujoco.rollout gets there first (Lesson 20.3).

Implementation

INPUT: cartpole.xml upright; a pendulum with the mujoco.pid plugin; c/minimal_step.c
PROCESS: mjd_transitionFD and its check against real steps; discrete LQR; a control callback's cost and its effect on mjd_transitionFD; the shipped plugins; compiling and running the C program against the package's library
OUTPUT: printed tables

```python file=examples/l20_4_extending.py “"”Lesson 20.4: derivatives, callbacks, plugins and the C API (Level 20).

INPUT cartpole.xml (2 dof, a force on the cart), balanced upright at the origin; a pendulum driven by the mujoco.pid actuator plugin; c/minimal_step.c PROCESS (1) mjd_transitionFD: the one-step linearization (A, B) of the upright cart-pole; how well A dx predicts one step from perturbations of size 1e-1 to 1e-7; centred against forward differences; (2) a discrete LQR controller from (A, B): which initial pole angles it balances in 5 s; (3) a control callback: its cost per step from Python, how often mjd_transitionFD calls it, and what it does to the B it reports; (4) the plugins shipped with the Python package, and the PID plugin holding a pendulum at 0.5 rad; (5) the C program c/minimal_step.c compiled against the package’s headers and library, against Python OUTPUT printed tables

Run: python examples/l20_4_extending.py (part 5 needs a C compiler, cc) “””

import os import re import shutil import subprocess import tempfile import time from pathlib import Path

import mujoco import numpy as np from scipy.linalg import solve_discrete_are

from mjcourse import model_path

CODE = Path(file).resolve().parents[1]

def upright() -> tuple[mujoco.MjModel, mujoco.MjData]: model = mujoco.MjModel.from_xml_path(str(model_path(“cartpole”))) data = mujoco.MjData(model) mujoco.mj_forward(model, data) return model, data

def linearize(model, data, eps: float = 1e-6, centered: bool = True) -> tuple[np.ndarray, np.ndarray]: nx = 2 * model.nv + model.na a, b = np.zeros((nx, nx)), np.zeros((nx, model.nu)) mujoco.mjd_transitionFD(model, data, eps, centered, a, b, None, None) return a, b

def next_state(model, x: np.ndarray, u: float = 0.0) -> np.ndarray: “"”One step from state x = (qpos, qvel) of the cart-pole (its qpos is a vector space: no quaternions).””” data = mujoco.MjData(model) data.qpos[:], data.qvel[:] = x[:model.nq], x[model.nq:] data.ctrl[0] = u mujoco.mj_step(model, data) return np.r_[data.qpos, data.qvel]

def balance(model, gain: np.ndarray, angle: float, seconds: float = 5.0) -> tuple[bool, float]: data = mujoco.MjData(model) data.qpos[1] = angle peak = 0.0 while data.time < seconds: u = float(-gain @ np.r_[data.qpos, data.qvel]) peak = max(peak, abs(u)) data.ctrl[0] = u # the actuator clamps it to +-20 N mujoco.mj_step(model, data) return bool(abs(data.qpos[1]) < 0.01 and abs(data.qpos[0]) < 1.4), peak

if name == “main”: model, data = upright() print(f”MuJoCo {mujoco.version}; cart-pole upright, timestep {model.opt.timestep} s”) a, b = linearize(model, data) print(“(1) mjd_transitionFD, centred, eps 1e-6: A (4 x 4) and B (4 x 1), state (cart, pole, cart velocity, pole rate)”) for row_a, row_b in zip(a, b): print(“ “ + “ “.join(f”{v:+.6f}” for v in row_a) + “ | “ + f”{row_b[0]:+.6e}”) a_fwd, b_fwd = linearize(model, data, centered=False) moving = mujoco.MjData(model) moving.qpos[:], moving.qvel[:] = [0.2, 0.6], [0.5, -1.0] mujoco.mj_forward(model, moving) gaps = [np.abs(linearize(model, moving, e)[0] - linearize(model, moving, e, centered=False)[0]).max() for e in (1e-6, 1e-4)] print(f” forward against centred differences, largest gap in A: {np.abs(a_fwd - a).max():.1e} here (a symmetric point); “ f”at the moving state (0.2, 0.6, 0.5, -1.0): {gaps[0]:.1e} with eps 1e-6, {gaps[1]:.1e} with eps 1e-4”) rng = np.random.default_rng(0) direction = rng.normal(size=4) direction /= np.linalg.norm(direction) x0 = np.zeros(4) base = next_state(model, x0) cells = [] for s in (1e-1, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6, 1e-7): dx = s * direction actual, predicted = next_state(model, x0 + dx) - base, a @ dx cells.append(f”{s:.0e}: {np.linalg.norm(actual - predicted) / np.linalg.norm(predicted):.1e}”) print(“ one step from a perturbation of size s, relative error of A dx: “ + “, “.join(cells))

print("(2) discrete LQR from (A, B): Q = diag(1, 10, 1, 1), R = 0.1")
q, r = np.diag([1.0, 10.0, 1.0, 1.0]), np.array([[0.1]])
p = solve_discrete_are(a, b, q, r)
gain = np.linalg.solve(r + b.T @ p @ b, b.T @ p @ a)[0]
closed = np.abs(np.linalg.eigvals(a - np.outer(b[:, 0], gain))).max()
print(f"    gain K = [{', '.join(f'{k:.2f}' for k in gain)}]; largest closed-loop eigenvalue magnitude {closed:.5f}")
for angle in (0.1, 0.3, 0.5, 0.7, 0.9):
    ok, peak = balance(model, gain, angle)
    print(f"    start at {angle:.1f} rad: {'balanced' if ok else 'not balanced'} after 5 s, largest commanded force {peak:6.1f} N")

print("(3) a control callback")
calls = [0]

def callback(m, d):
    calls[0] += 1
    d.ctrl[0] = float(-gain @ np.r_[d.qpos, d.qvel])
rates = {}
for label, cb in (("no callback", None), ("Python callback", callback)):
    mujoco.set_mjcb_control(cb)
    d = mujoco.MjData(model)
    d.qpos[1] = 0.1
    t0, n = time.perf_counter(), 20000
    for _ in range(n):
        mujoco.mj_step(model, d)
    rates[label] = (time.perf_counter() - t0) / n
print(f"    microseconds per step: without {1e6 * rates['no callback']:.2f}, with a Python control callback "
      f"{1e6 * rates['Python callback']:.2f}")
mujoco.set_mjcb_control(callback)
calls[0] = 0
a_cb, b_cb = linearize(model, data)
mujoco.set_mjcb_control(None)
print(f"    mjd_transitionFD with the callback installed called it {calls[0]} times; |B| with it {np.linalg.norm(b_cb):.2e} "
      f"(without: {np.linalg.norm(b):.2e}); A with it equals A - B K to {np.abs(a_cb - (a - np.outer(b[:, 0], gain))).max():.1e}")

print("(4) plugins shipped with the Python package")
plugin_dir = Path(mujoco.__file__).parent / "plugin"
for lib in sorted(plugin_dir.glob("*.so")) + sorted(plugin_dir.glob("*.dll")) + sorted(plugin_dir.glob("*.dylib")):
    names = sorted(set(re.findall(rb"mujoco\.(?!so\b)[a-z_]+(?:\.[a-z_]+)*", lib.read_bytes())))  # not "libmujoco.so"
    print(f"    {lib.name}: " + ", ".join(n.decode() for n in names))
pid = mujoco.MjModel.from_xml_string("""<mujoco>

</mujoco>”””) d = mujoco.MjData(pid) d.ctrl[0] = 0.5 while d.time < 6.0: mujoco.mj_step(pid, d) print(f” mujoco.pid (kp 40, ki 20, kd 4) holding a 1 kg pendulum at 0.5 rad: angle after 6 s {d.qpos[0]:.4f} rad, “ f”integrator state (act) {d.act[0]:.4f}”)

print("(5) c/minimal_step.c")
compiler = shutil.which("cc") or shutil.which("gcc") or shutil.which("clang")
if compiler is None or os.name == "nt":
    print("    no C compiler found (or Windows); skipped")
else:
    mj = Path(mujoco.__file__).parent
    library = next(mj.glob("libmujoco.so*"), None) or next(mj.glob("libmujoco*.dylib"))
    with tempfile.TemporaryDirectory() as tmp:
        exe = Path(tmp) / "minimal_step"
        cmd = [compiler, str(CODE / "c" / "minimal_step.c"), f"-I{mj / 'include'}", str(library), f"-Wl,-rpath,{mj}", "-o", str(exe)]
        subprocess.run(cmd, check=True)
        out = subprocess.run([str(exe), str(model_path("cartpole"))], check=True, capture_output=True, text=True).stdout.strip()
    print(f"    {out}")
    m = mujoco.MjModel.from_xml_path(str(model_path("cartpole")))
    d = mujoco.MjData(m)
    mujoco.mj_resetDataKeyframe(m, d, 0)
    for _ in range(1000):
        mujoco.mj_step(m, d)
    c_qpos = np.array([float(v) for v in out.split("qpos =")[1].split()])
    print(f"    Python, same model and steps: qpos = {' '.join(f'{v:.17g}' for v in d.qpos)}; "
          f"largest difference {np.abs(c_qpos - d.qpos).max():.1e}") ``` ```c file=c/minimal_step.c // minimal_step.c: the smallest useful MuJoCo program in C (Lesson 20.4). // Loads a model, resets it to its first keyframe, steps it 1000 times and prints the state. // // Build against the headers and library that ship with the Python package: //   MJ=$(python -c "import mujoco, os; print(os.path.dirname(mujoco.__file__))") //   cc minimal_step.c -I"$MJ/include" -L"$MJ" -l:libmujoco.so.3.14.0 -Wl,-rpath,"$MJ" -o minimal_step //   ./minimal_step ../src/mjcourse/models/cartpole.xml

#include

#include <mujoco/mujoco.h>

int main(int argc, char** argv) { if (argc < 2) { fprintf(stderr, “usage: %s model.xml\n”, argv[0]); return 1; }

char error[1000] = “”; mjModel* m = mj_loadXML(argv[1], NULL, error, sizeof(error)); if (!m) { fprintf(stderr, “could not load %s: %s\n”, argv[1], error); return 1; } mjData* d = mj_makeData(m); if (m->nkey > 0) { mj_resetDataKeyframe(m, d, 0); }

for (int k = 0; k < 1000; k++) { mj_step(m, d); }

printf(“MuJoCo %s: after 1000 steps t = %.3f s, qpos =”, mj_versionString(), d->time); for (int i = 0; i < m->nq; i++) { printf(“ %.17g”, d->qpos[i]); } printf(“\n”);

mj_deleteData(d); mj_deleteModel(m); return 0; }

Output:

```text
MuJoCo 3.14.0; cart-pole upright, timestep 0.002 s
(1) mjd_transitionFD, centred, eps 1e-6: A (4 x 4) and B (4 x 1), state (cart, pole, cart velocity, pole rate)
    +1.000000 -0.000005 +0.002000 +0.000000   |   +3.799370e-06
    +0.000000 +1.000110 +0.000001 +0.002000   |   -9.333404e-06
    +0.000000 -0.002747 +0.999810 +0.000009   |   +1.899685e-03
    +0.000000 +0.054946 +0.000467 +0.999813   |   -4.666702e-03
    forward against centred differences, largest gap in A: 0.0e+00 here (a symmetric point); at the moving state (0.2, 0.6, 0.5, -1.0): 2.1e-08 with eps 1e-6, 2.1e-06 with eps 1e-4
    one step from a perturbation of size s, relative error of A dx: 1e-01: 1.2e-06, 1e-02: 1.2e-08, 1e-03: 1.2e-10, 1e-04: 1.2e-12, 1e-05: 8.5e-15, 1e-06: 3.2e-15, 1e-07: 3.4e-15
(2) discrete LQR from (A, B): Q = diag(1, 10, 1, 1), R = 0.1
    gain K = [-3.11, -45.58, -5.68, -9.18]; largest closed-loop eigenvalue magnitude 0.99734
    start at 0.1 rad: balanced after 5 s, largest commanded force    4.6 N
    start at 0.3 rad: balanced after 5 s, largest commanded force   13.7 N
    start at 0.5 rad: balanced after 5 s, largest commanded force   22.8 N
    start at 0.7 rad: balanced after 5 s, largest commanded force   31.9 N
    start at 0.9 rad: not balanced after 5 s, largest commanded force 2026.2 N
(3) a control callback
    microseconds per step: without 2.68, with a Python control callback 9.59
    mjd_transitionFD with the callback installed called it 11 times; |B| with it 0.00e+00 (without: 5.04e-03); A with it equals A - B K to 1.4e-13
(4) plugins shipped with the Python package
    libactuator.so: mujoco.pid
    libelasticity.so: mujoco.elasticity.cable
    libsdf_plugin.so: mujoco.sdf.bolt, mujoco.sdf.bowl, mujoco.sdf.gear, mujoco.sdf.nut, mujoco.sdf.torus
    libsensor.so: mujoco.sensor.touch_grid
    mujoco.pid (kp 40, ki 20, kd 4) holding a 1 kg pendulum at 0.5 rad: angle after 6 s 0.4997 rad, integrator state (act) 0.0583
(5) c/minimal_step.c
    MuJoCo 3.14.0: after 1000 steps t = 2.000 s, qpos = -0.050354155334510445 1.899214965103317
    Python, same model and steps: qpos = -0.050354155334510445 1.899214965103317; largest difference 0.0e+00

Debugging

B is zero or A looks closed-loop. A control callback is installed while you differentiate. Remove it.

A linearization predicts badly. The state is far from where you linearized, or eps is too large or too small. Check $A\,\delta x$ against a real step for a few sizes, as part (1) does.

A plugin is not found. The library is not loaded. The Python package loads its own on import; for your own plugins call mujoco.mj_loadPluginLibrary with its path first.

The C program does not link or cannot find the library at run time. Pass the library’s directory with -L or its full path, and -Wl,-rpath (or LD_LIBRARY_PATH) so the loader finds it when the program runs.

Exercise

Linearize the cart-pole at a pole angle of 0.3 rad (held there by a constant force you compute from the equations of motion) and compare its $A$ with the upright one. Build LQR from each and compare the range of starting angles each balances.

Challenge

Write a mjcb_passive callback that adds aerodynamic drag $-c\,v\lvert v\rvert$ to the cart, measure its cost per step from Python, then implement the same force as a passive-force plugin in C and compare speed and results.

Research connection

[!research] Derivatives in planning and learning Gradient-based trajectory optimization (iLQG; Tassa, Erez and Todorov, IROS 2012, in the reading list) uses exactly these one-step linearizations, and model-based reinforcement learning methods that differentiate through a simulator need them too. Contact makes them fragile. Repeat part (1) at a state where a contact switches on or off within the perturbation and watch the error stop shrinking as $s^2$ (inference; Lessons 18.2 and 20.2 show the same non-smoothness from other angles).

{"id": "20.4-check", "title": "Knowledge check", "questions": [
  {"kind": "mcq", "q": "Which entries of the upright cart-pole's A make the upright pose unstable?",
   "options": ["A[0,2] = 0.002 alone", "A[3,1] = 0.0549 together with A[1,3] = 0.002, which give A an eigenvalue of about 1.0105", "A[2,2] = 0.99981", "B[3]"],
   "answer": 1,
   "explain": "<p>Gravity feeds the pole angle into the pole rate, and the timestep feeds the rate back into the angle; the 2 x 2 block has an eigenvalue above 1.</p>"},
  {"kind": "mcq", "q": "With a control callback that sets ctrl installed, what did mjd_transitionFD return for B?",
   "options": ["The correct B", "Exactly zero", "Noise", "It raised an error"],
   "answer": 1,
   "explain": "<p>The callback overwrote every control perturbation; A came back as A - B K.</p>"},
  {"kind": "numeric", "q": "How many times did mjd_transitionFD (centred) call the callback for the cart-pole (nu 1, nv 2)?",
   "answer": 11, "tol": 0.01,
   "explain": "<p>1 nominal step + 2 per control + 2 x 2 per velocity + 2 x 2 per position.</p>"},
  {"kind": "mcq", "q": "What travels with the model and keeps its own state: a callback or a plugin?",
   "options": ["A callback", "A plugin", "Both", "Neither"],
   "answer": 1,
   "explain": "<p>Callbacks are process-wide function pointers; plugins are declared in MJCF and their state lives in mjData.</p>"}
]}

Next

Level 21 turns to research-grade environments: designing them, measuring with them, and reading other people’s benchmarks.