In [ ]:

from __future__ import annotations

import pathlib
from typing import Tuple

import jax.numpy as jnp
import matplotlib.pyplot as plt
import mujoco
from mujoco import MjData, MjModel, MjvOption, Renderer, mjtVisFlag
import numpy as np
from jaxonomy import library, simulate, SimulatorOptions
from jaxonomy.framework import DiagramBuilder

import io

from matplotlib.backends.backend_agg import FigureCanvasAgg
from PIL import Image

from __future__ import annotations import pathlib from typing import Tuple import jax.numpy as jnp import matplotlib.pyplot as plt import mujoco from mujoco import MjData, MjModel, MjvOption, Renderer, mjtVisFlag import numpy as np from jaxonomy import library, simulate, SimulatorOptions from jaxonomy.framework import DiagramBuilder import io from matplotlib.backends.backend_agg import FigureCanvasAgg from PIL import Image

Triple inverted pendulum: Jaxonomy LQR + MuJoCo kick-up and stabilization¶

This tutorial designs an LQR gain $K$ with Jaxonomy's library.LinearQuadraticRegulator (a continuous-time algebraic-Riccati solve on the MuJoCo-linearized plant), then validates it against MuJoCo physics on a triple pendulum on a cart. Jaxonomy is the control-design engine; MuJoCo mj_step is the high-fidelity reference. We simulate in two phases:

1. Kick-up (open loop): from the upright equilibrium, apply a short cosine burst on the cart force, $u(t)=A\cos(2\pi f t)$, to excite the chain (impulse-like start because $u(0)=A$).
2. Stabilization (closed loop): after a fixed time, switch to $u=-\mathrm{clip}(Kx)$ using the same $K$ as in the Jaxonomy diagram. $Q$ and $R$ are chosen so the gains are not immediately saturated (large $R$), which is critical for mj_step to track the continuous-time LQR design.

Visualization uses MuJoCo forward kinematics (mj_forward + body positions) and, when OpenGL is available, an optional MuJoCo Renderer pass. Animations are GIF via Pillow — no ffmpeg.

Reference vibe: cart–pendulum style demos on YouTube.

In [ ]:

# Marimo often runs cells without __file__; fall back to cwd (run from repo / docs/examples).
_here = globals().get("__file__")
HERE = pathlib.Path(_here).resolve().parent if _here else pathlib.Path.cwd()
xml_path = HERE / "mujoco" / "assets" / "triple_pendulum_cart.xml"
assert xml_path.is_file(), f"Missing MJCF: {xml_path}"
xml = str(xml_path)

# Marimo often runs cells without __file__; fall back to cwd (run from repo / docs/examples). _here = globals().get("__file__") HERE = pathlib.Path(_here).resolve().parent if _here else pathlib.Path.cwd() xml_path = HERE / "mujoco" / "assets" / "triple_pendulum_cart.xml" assert xml_path.is_file(), f"Missing MJCF: {xml_path}" xml = str(xml_path)

In [ ]:

def linearize_at_key(
    xml_file: str, key_name: str = "upright", eps: float = 1e-5
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    m = MjModel.from_xml_path(xml_file)
    d = MjData(m)
    kid = m.key(key_name).id
    mujoco.mj_resetDataKeyframe(m, d, kid)
    mujoco.mj_forward(m, d)
    nq, nv, nu = m.nq, m.nv, m.nu
    q0 = d.qpos.copy()
    v0 = d.qvel.copy()
    u0 = d.ctrl.copy()

    def rhs(q: np.ndarray, v: np.ndarray, u: np.ndarray) -> np.ndarray:
        d.qpos[:] = q
        d.qvel[:] = v
        d.ctrl[:] = u
        mujoco.mj_forward(m, d)
        return np.concatenate([d.qvel, d.qacc])

    nx = nq + nv
    A = np.zeros((nx, nx))
    B = np.zeros((nx, nu))
    for j in range(nx):
        dx = np.zeros(nx)
        dx[j] = eps
        dq, dv = dx[:nq], dx[nq:]
        A[:, j] = (rhs(q0 + dq, v0 + dv, u0) - rhs(q0 - dq, v0 - dv, u0)) / (2 * eps)
    for j in range(nu):
        du = np.zeros(nu)
        du[j] = eps
        B[:, j] = (rhs(q0, v0, u0 + du) - rhs(q0, v0, u0 - du)) / (2 * eps)
    return A, B, q0, v0, u0

A_np, B_np, q_eq, v_eq, u_eq = linearize_at_key(xml, "upright")
# Large R → small gains so LQR is not instantly saturated; Q weights balance cart vs links.
# Tuned with MuJoCo `mj_step` for: brief open-loop kick (see rollout) then this LQR.
Q_np = np.diag([80.0, 60.0, 60.0, 60.0, 15.0, 15.0, 15.0, 15.0])
R_np = np.array([[5000.0]])
# Synthesize the stabilizing gain with Jaxonomy's LQR (continuous-time
# algebraic-Riccati solve). This is the SAME K the LinearQuadraticRegulator
# block builds in the diagram below; reading it off the block and driving
# the MuJoCo rollout with it makes Jaxonomy the control-design engine and
# MuJoCo the high-fidelity physics reference it is validated against.
_lqr_gain = library.LinearQuadraticRegulator(
    jnp.asarray(A_np),
    jnp.asarray(B_np),
    jnp.asarray(Q_np),
    jnp.asarray(R_np),
    name="lqr_gain",
)
K_np = np.asarray(_lqr_gain.K).reshape(1, -1)

A_j = jnp.asarray(A_np)
B_j = jnp.asarray(B_np)
Q_j = jnp.asarray(Q_np)
R_j = jnp.asarray(R_np)
q_eq_j = jnp.asarray(q_eq)

def linearize_at_key( xml_file: str, key_name: str = "upright", eps: float = 1e-5 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: m = MjModel.from_xml_path(xml_file) d = MjData(m) kid = m.key(key_name).id mujoco.mj_resetDataKeyframe(m, d, kid) mujoco.mj_forward(m, d) nq, nv, nu = m.nq, m.nv, m.nu q0 = d.qpos.copy() v0 = d.qvel.copy() u0 = d.ctrl.copy() def rhs(q: np.ndarray, v: np.ndarray, u: np.ndarray) -> np.ndarray: d.qpos[:] = q d.qvel[:] = v d.ctrl[:] = u mujoco.mj_forward(m, d) return np.concatenate([d.qvel, d.qacc]) nx = nq + nv A = np.zeros((nx, nx)) B = np.zeros((nx, nu)) for j in range(nx): dx = np.zeros(nx) dx[j] = eps dq, dv = dx[:nq], dx[nq:] A[:, j] = (rhs(q0 + dq, v0 + dv, u0) - rhs(q0 - dq, v0 - dv, u0)) / (2 * eps) for j in range(nu): du = np.zeros(nu) du[j] = eps B[:, j] = (rhs(q0, v0, u0 + du) - rhs(q0, v0, u0 - du)) / (2 * eps) return A, B, q0, v0, u0 A_np, B_np, q_eq, v_eq, u_eq = linearize_at_key(xml, "upright") # Large R → small gains so LQR is not instantly saturated; Q weights balance cart vs links. # Tuned with MuJoCo `mj_step` for: brief open-loop kick (see rollout) then this LQR. Q_np = np.diag([80.0, 60.0, 60.0, 60.0, 15.0, 15.0, 15.0, 15.0]) R_np = np.array([[5000.0]]) # Synthesize the stabilizing gain with Jaxonomy's LQR (continuous-time # algebraic-Riccati solve). This is the SAME K the LinearQuadraticRegulator # block builds in the diagram below; reading it off the block and driving # the MuJoCo rollout with it makes Jaxonomy the control-design engine and # MuJoCo the high-fidelity physics reference it is validated against. _lqr_gain = library.LinearQuadraticRegulator( jnp.asarray(A_np), jnp.asarray(B_np), jnp.asarray(Q_np), jnp.asarray(R_np), name="lqr_gain", ) K_np = np.asarray(_lqr_gain.K).reshape(1, -1) A_j = jnp.asarray(A_np) B_j = jnp.asarray(B_np) Q_j = jnp.asarray(Q_np) R_j = jnp.asarray(R_np) q_eq_j = jnp.asarray(q_eq)

In [ ]:

builder = DiagramBuilder()
plant = builder.add(
    library.MJX(
        file_name=xml,
        dt=None,
        key_frame_0="upright",
        name="plant",
    )
)
mux = builder.add(library.Multiplexer(2, name="state_mux"))
lqr = builder.add(
    library.LinearQuadraticRegulator(A_j, B_j, Q_j, R_j, name="lqr"),
)
builder.connect(plant.output_ports[0], mux.input_ports[0])
builder.connect(plant.output_ports[1], mux.input_ports[1])
builder.connect(mux.output_ports[0], lqr.input_ports[0])
builder.connect(lqr.output_ports[0], plant.input_ports[0])
diagram = builder.build()

mo.md(
    r"""
    ## Jaxonomy diagram

    `Multiplexer` stacks `qpos` and `qvel` into the $8$-vector $x$ for LQR.
    The regulator implements $u = -K x$ (equilibrium at the upright keyframe).
    """
)
diagram.pprint()

builder = DiagramBuilder() plant = builder.add( library.MJX( file_name=xml, dt=None, key_frame_0="upright", name="plant", ) ) mux = builder.add(library.Multiplexer(2, name="state_mux")) lqr = builder.add( library.LinearQuadraticRegulator(A_j, B_j, Q_j, R_j, name="lqr"), ) builder.connect(plant.output_ports[0], mux.input_ports[0]) builder.connect(plant.output_ports[1], mux.input_ports[1]) builder.connect(mux.output_ports[0], lqr.input_ports[0]) builder.connect(lqr.output_ports[0], plant.input_ports[0]) diagram = builder.build() mo.md( r""" ## Jaxonomy diagram `Multiplexer` stacks `qpos` and `qvel` into the $8$-vector $x$ for LQR. The regulator implements $u = -K x$ (equilibrium at the upright keyframe). """ ) diagram.pprint()

In [ ]:

"""MuJoCo `mj_step`: cosine **kick-up** impulse from upright, then Jaxonomy-tuned LQR."""
m = MjModel.from_xml_path(xml)
_mjd = MjData(m)
mujoco.mj_resetDataKeyframe(m, _mjd, m.key("upright").id)

_dt = float(m.opt.timestep)
_T = 8.0
n_steps = int(_T / _dt)
q_hist = np.zeros((n_steps, m.nq))
v_hist = np.zeros((n_steps, m.nv))
u_hist = np.zeros((n_steps, m.nu))
kick_phase = np.zeros(n_steps, dtype=bool)

# Open-loop cart force u(t)=A cos(2π f t) for t∈[0,T_kick); then u=-clip(Kx).
_T_kick = 0.35
_A_kick, _f_kick = 40.0, 2.5
_u_lim = 300.0

for _sk in range(n_steps):
    _t = _sk * _dt
    q_hist[_sk] = _mjd.qpos.copy()
    v_hist[_sk] = _mjd.qvel.copy()
    _x = np.concatenate([_mjd.qpos, _mjd.qvel])
    if _t < _T_kick:
        _u = float(
            np.clip(_A_kick * np.cos(2.0 * np.pi * _f_kick * _t), -_u_lim, _u_lim)
        )
        kick_phase[_sk] = True
    else:
        _u = float(np.clip(-(K_np @ _x.reshape(-1, 1)).item(), -_u_lim, _u_lim))
        kick_phase[_sk] = False
    u_hist[_sk] = _u
    _mjd.ctrl[:] = _u
    mujoco.mj_step(m, _mjd)

t_series = np.arange(n_steps) * _dt

"""MuJoCo `mj_step`: cosine **kick-up** impulse from upright, then Jaxonomy-tuned LQR.""" m = MjModel.from_xml_path(xml) _mjd = MjData(m) mujoco.mj_resetDataKeyframe(m, _mjd, m.key("upright").id) _dt = float(m.opt.timestep) _T = 8.0 n_steps = int(_T / _dt) q_hist = np.zeros((n_steps, m.nq)) v_hist = np.zeros((n_steps, m.nv)) u_hist = np.zeros((n_steps, m.nu)) kick_phase = np.zeros(n_steps, dtype=bool) # Open-loop cart force u(t)=A cos(2π f t) for t∈[0,T_kick); then u=-clip(Kx). _T_kick = 0.35 _A_kick, _f_kick = 40.0, 2.5 _u_lim = 300.0 for _sk in range(n_steps): _t = _sk * _dt q_hist[_sk] = _mjd.qpos.copy() v_hist[_sk] = _mjd.qvel.copy() _x = np.concatenate([_mjd.qpos, _mjd.qvel]) if _t < _T_kick: _u = float( np.clip(_A_kick * np.cos(2.0 * np.pi * _f_kick * _t), -_u_lim, _u_lim) ) kick_phase[_sk] = True else: _u = float(np.clip(-(K_np @ _x.reshape(-1, 1)).item(), -_u_lim, _u_lim)) kick_phase[_sk] = False u_hist[_sk] = _u _mjd.ctrl[:] = _u mujoco.mj_step(m, _mjd) t_series = np.arange(n_steps) * _dt

In [ ]:

mo.md(r"## Time series (MuJoCo `mj_step`: kick-up then LQR)")

_fig, _axes = plt.subplots(3, 1, figsize=(9, 7), sharex=True)
_kick = kick_phase.astype(bool)
if np.any(_kick):
    _tk = t_series[_kick]
    _axes[0].axvspan(float(_tk.min()), float(_tk.max()), color="orange", alpha=0.12, label="open-loop kick")
labels = [r"$x_{\mathrm{cart}}$", r"$\theta_1$", r"$\theta_2$", r"$\theta_3$"]
for _qi in range(4):
    _axes[0].plot(t_series, q_hist[:, _qi], label=labels[_qi])
_axes[0].set_ylabel("qpos [rad or m]")
_axes[0].legend(loc="upper right", fontsize=8)
_axes[0].grid(True, alpha=0.3)
_axes[0].set_title("Generalized positions")

for _vi in range(4):
    _axes[1].plot(t_series, v_hist[:, _vi], label=f"v{_vi}")
_axes[1].set_ylabel("qvel")
_axes[1].legend(loc="upper right", fontsize=8)
_axes[1].grid(True, alpha=0.3)
_axes[1].set_title("Generalized velocities")

_axes[2].plot(t_series, u_hist[:, 0], color="C3")
_axes[2].set_ylabel("u [N]")
_axes[2].set_xlabel("time [s]")
_axes[2].grid(True, alpha=0.3)
_axes[2].set_title("Cart force (saturated)")
plt.tight_layout()
plt.show()

mo.md(r"## Time series (MuJoCo `mj_step`: kick-up then LQR)") _fig, _axes = plt.subplots(3, 1, figsize=(9, 7), sharex=True) _kick = kick_phase.astype(bool) if np.any(_kick): _tk = t_series[_kick] _axes[0].axvspan(float(_tk.min()), float(_tk.max()), color="orange", alpha=0.12, label="open-loop kick") labels = [r"$x_{\mathrm{cart}}$", r"$\theta_1$", r"$\theta_2$", r"$\theta_3$"] for _qi in range(4): _axes[0].plot(t_series, q_hist[:, _qi], label=labels[_qi]) _axes[0].set_ylabel("qpos [rad or m]") _axes[0].legend(loc="upper right", fontsize=8) _axes[0].grid(True, alpha=0.3) _axes[0].set_title("Generalized positions") for _vi in range(4): _axes[1].plot(t_series, v_hist[:, _vi], label=f"v{_vi}") _axes[1].set_ylabel("qvel") _axes[1].legend(loc="upper right", fontsize=8) _axes[1].grid(True, alpha=0.3) _axes[1].set_title("Generalized velocities") _axes[2].plot(t_series, u_hist[:, 0], color="C3") _axes[2].set_ylabel("u [N]") _axes[2].set_xlabel("time [s]") _axes[2].grid(True, alpha=0.3) _axes[2].set_title("Cart force (saturated)") plt.tight_layout() plt.show()

In [ ]:

mo.md(
    r"""
    ## Animation (MuJoCo + Pillow GIF, no ffmpeg)

    Frames use **MuJoCo `mj_forward`** body positions (side view in $xz$). If OpenGL works, a **3D
    `Renderer`** clip is appended below the schematic.
    """
)

_dt = float(t_series[1] - t_series[0]) if len(t_series) > 1 else 0.01
# Cap the frame count so the GIFs stay well under marimo's 10 MB output limit.
_TARGET_FRAMES = 100
_stride = max(1, len(q_hist) // _TARGET_FRAMES)

def _quantize(_rgb: np.ndarray) -> Image.Image:
    return Image.fromarray(_rgb).quantize(colors=128, method=Image.Quantize.MEDIANCUT)

_body_names = ("cart", "link1", "link2", "link3")
_bids = tuple(
    mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, _n) for _n in _body_names
)

def _stick_rgb(_q: np.ndarray) -> np.ndarray:
    _d = MjData(m)
    _d.qpos[:] = _q
    _d.qvel[:] = 0.0
    mujoco.mj_forward(m, _d)
    _pts = np.stack([_d.xpos[_bi].copy() for _bi in _bids], axis=0)
    _w_in, _h_in, _dpi = 6.4, 4.0, 100
    _fig, _ax = plt.subplots(figsize=(_w_in, _h_in), dpi=_dpi)
    _ax.set_aspect("equal")
    _ax.plot(_pts[:, 0], _pts[:, 2], "o-", color="0.2", lw=3, markersize=9)
    _ax.fill_between([-2.6, 2.6], -0.05, 0.0, color="0.85", zorder=0)
    _ax.set_xlim(-2.5, 2.5)
    _ax.set_ylim(-0.1, 0.55)
    _ax.set_xlabel("x [m]")
    _ax.set_ylabel("z [m]")
    _ax.set_title("MuJoCo FK — triple pendulum on cart")
    _ax.grid(True, alpha=0.25)
    _canvas = FigureCanvasAgg(_fig)
    _canvas.draw()
    _rgba = np.asarray(_canvas.buffer_rgba())
    plt.close(_fig)
    return _rgba[:, :, :3].copy()

_pil_frames = [_quantize(_stick_rgb(q_hist[_fi]))
               for _fi in range(0, len(q_hist), _stride)]

_duration_ms = max(30, int(1000.0 * _stride * _dt))
_buf = io.BytesIO()
_pil_frames[0].save(
    _buf,
    format="GIF",
    save_all=True,
    append_images=_pil_frames[1:],
    duration=_duration_ms,
    loop=0,
    optimize=True,
)
_viz_rows = [
    mo.md(f"**Side view (matplotlib)** — {len(_pil_frames)} frames, "
          f"{len(_buf.getvalue()) / 1024:.0f} KB"),
    mo.image(_buf.getvalue()),
]
_gl_frames = []
try:
    _rw, _rh = 480, 270
    _gl_stride = _stride * 2
    _renderer = Renderer(m, width=_rw, height=_rh)
    _scene_option = MjvOption()
    _scene_option.flags[mjtVisFlag.mjVIS_JOINT] = True
    _mjd_gl = MjData(m)
    for _fi in range(0, len(q_hist), _gl_stride):
        _mjd_gl.qpos[:] = q_hist[_fi]
        _mjd_gl.qvel[:] = 0.0
        mujoco.mj_forward(m, _mjd_gl)
        _renderer.update_scene(_mjd_gl, scene_option=_scene_option)
        _gl_frames.append(_quantize(_renderer.render()))
    if _gl_frames:
        _buf2 = io.BytesIO()
        _gl_frames[0].save(
            _buf2,
            format="GIF",
            save_all=True,
            append_images=_gl_frames[1:],
            duration=_duration_ms * 2,
            loop=0,
            optimize=True,
        )
        _viz_rows.append(mo.md(f"**MuJoCo Renderer (3D)** — {len(_gl_frames)} frames, "
                               f"{len(_buf2.getvalue()) / 1024:.0f} KB"))
        _viz_rows.append(mo.image(_buf2.getvalue()))
except Exception as _e:
    _viz_rows.append(
        mo.md(
            f"*MuJoCo Renderer skipped ({type(_e).__name__}: {_e}). Schematic GIF above is enough for the tutorial.*"
        )
    )
mo.vstack(_viz_rows)

mo.md( r""" ## Animation (MuJoCo + Pillow GIF, no ffmpeg) Frames use **MuJoCo `mj_forward`** body positions (side view in $xz$). If OpenGL works, a **3D `Renderer`** clip is appended below the schematic. """ ) _dt = float(t_series[1] - t_series[0]) if len(t_series) > 1 else 0.01 # Cap the frame count so the GIFs stay well under marimo's 10 MB output limit. _TARGET_FRAMES = 100 _stride = max(1, len(q_hist) // _TARGET_FRAMES) def _quantize(_rgb: np.ndarray) -> Image.Image: return Image.fromarray(_rgb).quantize(colors=128, method=Image.Quantize.MEDIANCUT) _body_names = ("cart", "link1", "link2", "link3") _bids = tuple( mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, _n) for _n in _body_names ) def _stick_rgb(_q: np.ndarray) -> np.ndarray: _d = MjData(m) _d.qpos[:] = _q _d.qvel[:] = 0.0 mujoco.mj_forward(m, _d) _pts = np.stack([_d.xpos[_bi].copy() for _bi in _bids], axis=0) _w_in, _h_in, _dpi = 6.4, 4.0, 100 _fig, _ax = plt.subplots(figsize=(_w_in, _h_in), dpi=_dpi) _ax.set_aspect("equal") _ax.plot(_pts[:, 0], _pts[:, 2], "o-", color="0.2", lw=3, markersize=9) _ax.fill_between([-2.6, 2.6], -0.05, 0.0, color="0.85", zorder=0) _ax.set_xlim(-2.5, 2.5) _ax.set_ylim(-0.1, 0.55) _ax.set_xlabel("x [m]") _ax.set_ylabel("z [m]") _ax.set_title("MuJoCo FK — triple pendulum on cart") _ax.grid(True, alpha=0.25) _canvas = FigureCanvasAgg(_fig) _canvas.draw() _rgba = np.asarray(_canvas.buffer_rgba()) plt.close(_fig) return _rgba[:, :, :3].copy() _pil_frames = [_quantize(_stick_rgb(q_hist[_fi])) for _fi in range(0, len(q_hist), _stride)] _duration_ms = max(30, int(1000.0 * _stride * _dt)) _buf = io.BytesIO() _pil_frames[0].save( _buf, format="GIF", save_all=True, append_images=_pil_frames[1:], duration=_duration_ms, loop=0, optimize=True, ) _viz_rows = [ mo.md(f"**Side view (matplotlib)** — {len(_pil_frames)} frames, " f"{len(_buf.getvalue()) / 1024:.0f} KB"), mo.image(_buf.getvalue()), ] _gl_frames = [] try: _rw, _rh = 480, 270 _gl_stride = _stride * 2 _renderer = Renderer(m, width=_rw, height=_rh) _scene_option = MjvOption() _scene_option.flags[mjtVisFlag.mjVIS_JOINT] = True _mjd_gl = MjData(m) for _fi in range(0, len(q_hist), _gl_stride): _mjd_gl.qpos[:] = q_hist[_fi] _mjd_gl.qvel[:] = 0.0 mujoco.mj_forward(m, _mjd_gl) _renderer.update_scene(_mjd_gl, scene_option=_scene_option) _gl_frames.append(_quantize(_renderer.render())) if _gl_frames: _buf2 = io.BytesIO() _gl_frames[0].save( _buf2, format="GIF", save_all=True, append_images=_gl_frames[1:], duration=_duration_ms * 2, loop=0, optimize=True, ) _viz_rows.append(mo.md(f"**MuJoCo Renderer (3D)** — {len(_gl_frames)} frames, " f"{len(_buf2.getvalue()) / 1024:.0f} KB")) _viz_rows.append(mo.image(_buf2.getvalue())) except Exception as _e: _viz_rows.append( mo.md( f"*MuJoCo Renderer skipped ({type(_e).__name__}: {_e}). Schematic GIF above is enough for the tutorial.*" ) ) mo.vstack(_viz_rows)

In [ ]:

mo.md(
    r"""
    ## Optional: `jaxonomy.simulate` with MJX (same block diagram)

    This runs the **adaptive** ODE solver on the **MJX** continuous-time dynamics. For large tilts it can
    disagree with `mj_step` or drift from the upright setpoint; here we use a **small** initial angle on
    link 1 so the trajectory stays in the linear regulator’s useful neighborhood.
    """
)

_ctx = diagram.create_context()
_x0 = jnp.concatenate([q_eq_j.at[1].set(0.01), jnp.zeros(4)])
_ctx = _ctx.with_continuous_state([_x0])
_sim_options = SimulatorOptions(
    max_major_steps=4000,
    atol=1e-8,
    rtol=1e-7,
    max_minor_step_size=0.002,
)
_sim_results = simulate(
    diagram,
    _ctx,
    (0.0, 1.5),
    options=_sim_options,
    recorded_signals={"qpos": plant.output_ports[0]},
)
_t_j = np.asarray(_sim_results.time)
_q_j = np.asarray(_sim_results.outputs["qpos"])
_opt_fig, _opt_ax = plt.subplots(1, 1, figsize=(9, 3))
_opt_ax.plot(_t_j, _q_j[:, 0], label=r"$x_{\mathrm{cart}}$")
_opt_ax.plot(_t_j, _q_j[:, 1], label=r"$\theta_1$")
_opt_ax.set_xlabel("time [s]")
_opt_ax.set_ylabel("qpos")
_opt_ax.grid(True, alpha=0.3)
_opt_ax.legend()
_opt_ax.set_title("MJX + jaxonomy.simulate (optional; small IC)")
plt.tight_layout()
plt.show()

mo.md( r""" ## Optional: `jaxonomy.simulate` with MJX (same block diagram) This runs the **adaptive** ODE solver on the **MJX** continuous-time dynamics. For large tilts it can disagree with `mj_step` or drift from the upright setpoint; here we use a **small** initial angle on link 1 so the trajectory stays in the linear regulator’s useful neighborhood. """ ) _ctx = diagram.create_context() _x0 = jnp.concatenate([q_eq_j.at[1].set(0.01), jnp.zeros(4)]) _ctx = _ctx.with_continuous_state([_x0]) _sim_options = SimulatorOptions( max_major_steps=4000, atol=1e-8, rtol=1e-7, max_minor_step_size=0.002, ) _sim_results = simulate( diagram, _ctx, (0.0, 1.5), options=_sim_options, recorded_signals={"qpos": plant.output_ports[0]}, ) _t_j = np.asarray(_sim_results.time) _q_j = np.asarray(_sim_results.outputs["qpos"]) _opt_fig, _opt_ax = plt.subplots(1, 1, figsize=(9, 3)) _opt_ax.plot(_t_j, _q_j[:, 0], label=r"$x_{\mathrm{cart}}$") _opt_ax.plot(_t_j, _q_j[:, 1], label=r"$\theta_1$") _opt_ax.set_xlabel("time [s]") _opt_ax.set_ylabel("qpos") _opt_ax.grid(True, alpha=0.3) _opt_ax.legend() _opt_ax.set_title("MJX + jaxonomy.simulate (optional; small IC)") plt.tight_layout() plt.show()

Summary¶

• Model: four generalized positions (cart slide + three hinges) and matching velocities.
• Linearization: central differences on $[ \dot{q}; \dot{v} ] = [ v; a(q,v,u) ]$ at the upright keyframe.
• LQR: Jaxonomy's LinearQuadraticRegulator yields $K$ (continuous-time Riccati solve) and applies $u = -K x$ in the diagram; the MuJoCo mj_step rollout is driven by that same $K$.
• Kick + LQR: short open-loop cosine cart pulse, then saturated linear feedback $K$ from LQR.
• Visualization: schematic GIF (Pillow) from MuJoCo FK; optional 3D GIF from Renderer when GL works.

Full swing-up from hanging needs a dedicated nonlinear / MPC policy; this notebook focuses on a reliable disturbance + recovery demo with the same $K$ wired in Jaxonomy.