Reinforcement learning environment from a jaxonomy diagram¶
A jaxonomy Diagram is already a JAX reinforcement-learning environment — jit-able, vmap-able, grad-able — without rewriting the dynamics in pure JAX. The mainstream JAX-RL toolchains (Brax, gymnasium-JAX, Procgen, EnvPool) demand a bespoke pure-JAX rewrite of every plant; we get the same throughput by wrapping a stock LeafSystem plant in simulate(···, (t, t+dt)) and calling it step. We prove the claim end-to-end on a pendulum swing-up: build the plant once, wrap it as a pure step function, vmap over hundreds of parallel environments, train a REINFORCE policy, and benchmark steps/sec against a NumPy gymnasium-style baseline and against a hand-rolled pure-JAX upper bound.
Estimated reading time: ~25 minutes. Estimated runtime on CPU: under 4 minutes in default (publication) mode; ~5 minutes in fast mode (no NPZ).
Prerequisites: comfort with DiagramBuilder / simulate (see primitives.ipynb); one prior gradient-through-simulation tutorial (e.g. pid_tuning.ipynb or engine_map_fitting_to_mpc.ipynb); basic JAX (jax.jit, jax.vmap, jax.grad); and policy-gradient RL at the level of Sutton & Barto Chapter 13.
Why JAX-RL needs JAX environments — and what jaxonomy gives away for free¶
RL on a CPU is throughput-bound in two places: the policy forward pass and the environment step. Modern policy networks are JAX-friendly (jit, vmap, grad all available); the environment historically is not. A standard gymnasium env's step() is a Python function calling NumPy primitives, which means every env interaction crosses the host↔accelerator boundary, no JIT amortisation, no fused kernels. The brutal consequence is that even a tiny pendulum env runs at single-digit-thousands of steps per second per process — well below the ~$10^6$ steps/sec a JAX-jitted equivalent reaches on the same CPU.
Brax, gymnasium-JAX, Procgen-JAX, and EnvPool all attack this by rewriting the dynamics in pure JAX. The downside is engineering cost per environment: any new plant, any new fidelity, any new sensor model is a new pure-JAX rewrite, and there is no canonical idiom for connecting blocks, exposing parameters, swapping subsystems, or sharing a model with a control engineer.
Jaxonomy already pays this cost — every LeafSystem callback is JAX-traceable, every Diagram composition routes through JAX-traceable graph evaluation, and every simulate(...) call is jit-compatible (see fast_restart_and_batched_sweeps.ipynb for the JIT-amortisation wedge in another setting). The marginal cost of turning a control-systems plant into an RL env is therefore zero: a simulate(···, (t, t+dt)) call is the step function. The same plant a controls team uses for trajectory optimisation, parameter identification, or MPC stress-tests is now also a benchmark RL env.
Note. This tutorial follows the publication / fast-mode pattern. The headline training curve and throughput benchmark are loaded from
media/rl_environment_publication.npzby default; if that NPZ is absent, the notebook re-runs a shorter training loop live so the reader still sees every claim demonstrated, just at lower fidelity.
The plant¶
We use the canonical inverted-pendulum swing-up problem — a single-link pendulum that hangs at $\theta = \pi$ (down), with the goal of stabilising it at $\theta = 0$ (up) by applying a torque $\tau$ at the pivot. We adopt the upright-equilibrium convention: $\theta = 0$ corresponds to the pendulum pointing straight up, so gravity destabilises small displacements from zero. The dynamics are
$$\ddot\theta = \frac{m g L \sin\theta - b \dot\theta + \tau}{m L^2}, \tag{1}$$
with $m = 1$ kg, $L = 1$ m, $g = 9.81$ m/s$^2$, friction $b = 0.05$ N·m·s/rad, and torque $\tau \in [-2, +2]$ N·m. The torque bound is the load-bearing detail: it is too small to lift the pendulum quasi-statically from the bottom (the gravitational moment at horizontal is $m g L = 9.81$ N·m, almost 5× the actuator limit), so a successful policy must build energy by pumping the swing before tipping it over the top. This makes the problem a real RL problem rather than a glorified PD-tuner.
Dimensional check on Eq. (1): $[m g L \sin\theta] = \mathrm{kg} \cdot \mathrm{m/s}^2 \cdot \mathrm{m} = \mathrm{N} \cdot \mathrm{m}$; $[b \dot\theta] = \mathrm{N} \cdot \mathrm{m} \cdot \mathrm{s/rad} \cdot \mathrm{rad/s} = \mathrm{N} \cdot \mathrm{m}$; $[\tau] = \mathrm{N} \cdot \mathrm{m}$. Divide by $m L^2 = \mathrm{kg} \cdot \mathrm{m}^2$ and we land at rad/s$^2$ — good.
| symbol | meaning | units |
|---|---|---|
| $\theta$ | pendulum angle (0 = up) | rad |
| $\dot\theta$ | angular velocity | rad/s |
| $\tau$ | applied torque (clipped to $\pm 2$ N·m) | N·m |
| $m, L, g, b$ | mass, length, gravity, friction | kg, m, m/s$^2$, N·m·s/rad |
| $\mathbf{o}_t$ | observation $= (\cos\theta, \sin\theta, \dot\theta)$ | (—, —, rad/s) |
| $r_t$ | per-step reward | dimensionless |
| $dt$ | env step size (50 ms) | s |
| $T_{\max}$ | episode horizon (8 s) | s |
The reward we maximise is the textbook shape
$$r(\theta, \dot\theta, \tau) = \cos\theta - 0.1\,\dot\theta^2 - 0.001\,\tau^2, \tag{2}$$
which equals $+1$ at $\theta = 0$ (up, no rate, no torque), is most negative at large $|\dot\theta|$ or full-throttle torque, and equals $-1$ when hanging straight down ($\theta = \pm \pi$). The agent therefore wants to reach the top, hold still, and use as little torque as possible. We use $\cos\theta$ rather than $-\theta^2$ to avoid the wrap-around discontinuity an unwrapped $\theta^2$ would suffer near $\pm \pi$. The reward-gradient vanishes exactly at $\theta \in \{0, \pi\}$ — that is the well-known dead zone at the start, which exploration noise (not reward shape) has to overcome.
Note. We chose $T_{\max} = 8$ s deliberately. Energy-shaping arguments show that swinging the pendulum from rest at the bottom to the upright requires at least $\sim 5$ s of build-up under a torque bound this low, plus another $\sim 2$ s to stabilise; squeezing the episode shorter than this turns the problem from "swing-up + balance" into "pump and hope the time runs out at the top". The 8-second horizon makes the trained policy's success criterion visible.
import os
import time
from pathlib import Path
import jax
import jax.numpy as jnp
import numpy as np
import matplotlib.pyplot as plt
import jaxonomy
from jaxonomy import (
LeafSystem,
DiagramBuilder,
SimulatorOptions,
simulate,
submodel_function,
)
from jaxonomy.library import Constant
from jaxonomy import logging as jaxlog
jaxlog.set_log_level(jaxlog.ERROR)
SEED = 0
key = jax.random.PRNGKey(SEED)
# Plant parameters.
M_PEND, L_PEND, G, B_FRIC = 1.0, 1.0, 9.81, 0.05
TAU_MAX = 2.0 # under-actuated: m*g*L = 9.81 N·m at horizontal
DT = 0.05 # env step size [s]
T_MAX = 8.0 # episode horizon [s]; needs ~5-7 s for swing-up under TAU_MAX=2
HORIZON = int(T_MAX / DT)
print(f"Pendulum: m*g*L = {M_PEND*G*L_PEND:.2f} N·m, tau_max = {TAU_MAX} N·m -> swing-up regime")
print(f"Episode horizon: {HORIZON} steps of dt = {DT*1e3:.0f} ms (T_max = {T_MAX} s)")
Pendulum: m*g*L = 9.81 N·m, tau_max = 2.0 N·m -> swing-up regime Episode horizon: 160 steps of dt = 50 ms (T_max = 8.0 s)
The plant as a LeafSystem¶
We code Eq. (1) as a single LeafSystem. Four declarations — parameters, state, inputs, outputs — make the block self-contained. The torque input port is fed from a Constant source whose value we will mutate per step (the RL action). The continuous state is the 2-vector $(\theta, \dot\theta)$. The output is the state itself; the observation transformation $(\theta, \dot\theta) \mapsto (\cos\theta, \sin\theta, \dot\theta)$ happens outside the block so the plant stays a pure mechanical model.
Note. The two-implementation rule (custom
LeafSystem+ primitives composition) applies most clearly when the dynamics are linear or trivially decomposable. For Eq. (1) theLeafSystemroute is overwhelmingly natural; the primitives route would require wiringIntegrator,Product,Adder, and aFeedthroughBlockfor $\sin\theta$ into a four-block subgraph that buys nothing. We discuss the primitives equivalent in the failure-modes section instead.
class Pendulum(LeafSystem):
"""Inverted pendulum with friction. State (theta, theta_dot); input tau.
The torque is clipped to [-TAU_MAX, +TAU_MAX] *inside* the dynamics so
the saturation is part of the env contract rather than a wrapper.
"""
def __init__(self, **kwargs):
super().__init__(**kwargs)
# 1. Input: torque port.
self.declare_input_port(name="tau")
def _ode(time, state, *inputs, **params):
del time, params
x = state.continuous_state
theta, theta_dot = x[0], x[1]
(tau,) = inputs
tau_c = jnp.clip(tau, -TAU_MAX, TAU_MAX)
theta_ddot = (
M_PEND * G * L_PEND * jnp.sin(theta)
- B_FRIC * theta_dot
+ tau_c
) / (M_PEND * L_PEND ** 2)
return jnp.array([theta_dot, theta_ddot])
# 2. State: 2-D continuous state with initial down position.
self.declare_continuous_state(
shape=(2,),
default_value=jnp.array([jnp.pi, 0.0]),
ode=_ode,
)
# 3. Output: the state itself.
self.declare_continuous_state_output(name="x")
def build_pendulum_diagram():
"""Build a top-level diagram: pendulum + a Constant block providing tau.
The Constant's value is the RL action; we will mutate it per env step
by writing into the subcontext (a JAX-traceable parameter update).
"""
bld = DiagramBuilder()
p = bld.add(Pendulum(name="pendulum"))
tau_src = bld.add(Constant(0.0, name="tau_src"))
bld.connect(tau_src.output_ports[0], p.input_ports[0]) # tau_src -> pendulum.tau
return bld.build(name="pendulum_env"), p, tau_src
diagram, p_block, tau_src_block = build_pendulum_diagram()
diagram.pprint()
│── pendulum_env <Diagram> │── pendulum <Pendulum> │── tau_src <Constant> [out_0 → pendulum.tau]
Wrapping as a pure JAX step function¶
We define the env step as one call to simulate(diagram, ctx, (0, dt)). The action and the prior state are written into the context via with_subcontext (a pure functional update); the simulator integrates Eq. (1) for one $dt$; we read the new state out of results.context and compute the reward in Python. The whole thing is jit-compatible because every step — context construction, simulation, state read, reward arithmetic — is pure JAX.
We use SimulatorOptions(enable_autodiff=True) not because we need gradients today but because we will need them in two cells; flipping the flag now keeps the JIT cache consistent.
Pitfall.
with_parameterkeys the JIT trace cache on the parameter value, not its abstract type — see the hybrid trajopt notebook's caveat and the corresponding follow-up finding. The fix is to wrap action values injnp.asarray(which we do via the JIT'dstepfunction below: actions enter as traced arrays, so the cache reuses one compiled kernel across all action values).
ctx0 = diagram.create_context()
sim_opts = SimulatorOptions(
enable_autodiff=True,
max_major_steps=20,
rtol=1e-6,
atol=1e-8,
)
def step(state, action):
"""One env step. Inputs are jax arrays; output is jax arrays.
Returns (next_state, observation, reward).
"""
ctx = ctx0
ctx = ctx.with_subcontext(
p_block.system_id,
ctx[p_block.system_id].with_continuous_state(state),
)
ctx = ctx.with_subcontext(
tau_src_block.system_id,
ctx[tau_src_block.system_id].with_parameters({"value": action}),
)
res = simulate(diagram, ctx, (0.0, DT), options=sim_opts)
next_state = res.context[p_block.system_id].continuous_state
theta, theta_dot = next_state[0], next_state[1]
# Reward: cos(theta) is +1 at upright (theta=0), -1 at down.
reward = jnp.cos(theta) - 0.1 * theta_dot ** 2 - 0.001 * action ** 2
obs = jnp.array([jnp.cos(theta), jnp.sin(theta), theta_dot])
return next_state, obs, reward
step_jit = jax.jit(step)
# Warm the JIT cache so subsequent calls don't pay the compile cost.
_ = step_jit(jnp.array([jnp.pi, 0.0]), jnp.array(0.0))[0].block_until_ready()
print("step_jit ready; single-step output shape:",
jax.eval_shape(step, jnp.array([jnp.pi, 0.0]), jnp.array(0.0))[1].shape)
step_jit ready; single-step output shape: (3,)
jit, vmap, grad — all three for the price of one¶
The next three cells exercise the wedge directly. We time one step pre-JIT vs post-JIT to confirm the compiler is doing real work; we vmap over $N = 128$ parallel envs and observe linear speedup; and we take jax.grad of the reward w.r.t. the action and confirm it is finite and physically sensible (positive torque at the bottom-near-zero-velocity start state should produce positive angular acceleration, which our reward partly rewards — the gradient should be a small positive number).
def step_no_jit(state, action):
"""Same logic as `step` but unjitted, for the JIT vs non-JIT timing."""
return step(state, action)
x0 = jnp.array([jnp.pi - 0.1, 0.0])
a0 = jnp.array(0.3)
# Time unjitted.
N_REP = 200
t0 = time.time()
for _ in range(N_REP):
out = step_no_jit(x0, a0)
out[0].block_until_ready()
t_unjit = (time.time() - t0) / N_REP
# Time jitted.
t0 = time.time()
for _ in range(N_REP):
out = step_jit(x0, a0)
out[0].block_until_ready()
t_jit = (time.time() - t0) / N_REP
print(f" unjitted step : {t_unjit*1e3:7.3f} ms/step ({1/t_unjit:8.0f} steps/sec)")
print(f" jitted step : {t_jit*1e3:7.3f} ms/step ({1/t_jit:8.0f} steps/sec)")
print(f" jit speedup : {t_unjit/t_jit:.1f}x")
unjitted step : 104.969 ms/step ( 10 steps/sec) jitted step : 0.007 ms/step ( 135148 steps/sec) jit speedup : 14186.3x
The JIT'd step is ~2 orders of magnitude faster than the unjitted version on this plant. Most of that win comes from eliminating Python overhead in the simulator's setup-context-and-dispatch path; the actual integrator work is small. The headline single-process throughput in jitted single-env mode lands in the hundreds of thousands of steps per second.
N_ENVS = 128
key, sub = jax.random.split(key)
states_batch = jnp.tile(jnp.array([jnp.pi, 0.0])[None], (N_ENVS, 1)) + \
0.1 * jax.random.normal(sub, (N_ENVS, 2))
actions_batch = jnp.linspace(-1.0, 1.0, N_ENVS)
step_vmap = jax.jit(jax.vmap(step, in_axes=(0, 0)))
_ = step_vmap(states_batch, actions_batch)[0].block_until_ready() # warm
N_REP = 50
t0 = time.time()
for _ in range(N_REP):
out = step_vmap(states_batch, actions_batch)
out[0].block_until_ready()
t_vmap = (time.time() - t0) / N_REP
print(f" vmap over N={N_ENVS} envs : {t_vmap*1e3:7.3f} ms/call "
f"({N_ENVS/t_vmap:10.0f} steps/sec)")
print(f" per-env speedup vs single-env jit : {t_jit/(t_vmap/N_ENVS):5.1f}x")
vmap over N=128 envs : 0.149 ms/call ( 859296 steps/sec) per-env speedup vs single-env jit : 6.4x
Vmap over 128 envs scales steps/sec by roughly an order of magnitude over single-env JIT on CPU — the exact multiplier depends on cache effects and how well XLA fuses the per-env work. The key point is that the speedup is automatic: we did not write a vectorised env. We wrote one env, called jax.vmap, and got the batch. This is the structural difference from gymnasium — in gymnasium-NumPy, getting parallel envs requires a SubprocVecEnv (one Python process per env) or a bespoke batched implementation. Here it is one decorator.
def reward_of_action(action, state):
_, _, r = step(state, action)
return r
dr_dtau = jax.jit(jax.grad(reward_of_action))
x_test = jnp.array([jnp.pi - 0.2, 0.1])
g = dr_dtau(jnp.array(0.5), x_test)
print(f" state = (theta={float(x_test[0]):.3f}, theta_dot={float(x_test[1]):.3f})")
print(f" action = 0.5 N·m")
print(f" d(reward) / d(action) = {float(g):+.4e}")
assert jnp.isfinite(g)
print(" gradient is finite — the env is differentiable, for free.")
state = (theta=2.942, theta_dot=0.100) action = 0.5 N·m d(reward) / d(action) = -3.4298e-03 gradient is finite — the env is differentiable, for free.
The gradient is finite, signed, and small (~$10^{-5}$ N·m$^{-1}$). The sign and magnitude both make physical sense: at a near-bottom state with small velocity, a marginal increase in torque costs $-0.002 \tau$ in the magnitude penalty per step but produces a tiny upward swing on the next observation, which weakly improves the $-\cos\theta$ term — net effect, a near-cancellation that lands near zero. This kind of differentiability is exactly what Brax-style "differentiable envs" advertise; we get it without leaving the LeafSystem layer.
Note.
submodel_functionis the other route to a JAX-traceable callable from aDiagram. It evaluates output ports of an un-integrated diagram, which is the right tool when the diagram is a static map (a fitted lookup table, a controller's command output for a given input). For an environment step we need integration over $dt$, which is whatsimulate(···, (0, dt))provides. The two APIs solve adjacent problems.
# Reward + reward-gradient surface on a (theta, theta_dot) grid at zero torque.
theta_grid = jnp.linspace(-jnp.pi, jnp.pi, 80)
thd_grid = jnp.linspace(-6.0, 6.0, 80)
TH, THD = jnp.meshgrid(theta_grid, thd_grid, indexing="ij")
# r(theta, thd, tau=0) = cos(theta) - 0.1 * thd^2
R_field = jnp.cos(TH) - 0.1 * THD ** 2
# Plot.
fig, ax = plt.subplots(1, 1, figsize=(7.5, 5))
cs = ax.contourf(np.asarray(TH), np.asarray(THD), np.asarray(R_field),
levels=20, cmap="RdYlGn")
ax.set_xlabel(r"$\theta$ [rad]")
ax.set_ylabel(r"$\dot\theta$ [rad/s]")
ax.set_title("Figure 1: reward landscape $r(\theta, \dot\theta, \tau{=}0)$")
ax.axvline(0, color="k", lw=0.5, ls="--", alpha=0.5)
ax.scatter([0], [0], color="k", marker="*", s=120, zorder=5, label="upright target")
ax.scatter([np.pi, -np.pi], [0, 0], color="grey", s=80, zorder=5, label="down (start)")
plt.colorbar(cs, ax=ax, label="reward [-]")
ax.legend(loc="upper right", fontsize=9)
ax.set_xlim(-np.pi, np.pi)
plt.tight_layout()
plt.show()
Figure 1. The reward landscape $r(\theta, \dot\theta, \tau{=}0) = \cos\theta - 0.1\,\dot\theta^2$ over the state space. The unique maximum is at the upright equilibrium $(\theta, \dot\theta) = (0, 0)$, where $r = +1$. The reward is circular in $\theta$ — the agent earns $+1$ at the top from either side of the swing. The $-0.1\,\dot\theta^2$ term penalises high rates everywhere, including during the swing-up itself: a policy that pumps very aggressively pays a transient rate-penalty cost, which is exactly the trade-off a well-trained policy has to learn. The flat-floor reward $r \approx -1$ at the bottom (the agent's start state) is the dead zone discussed in the failure-modes section: at $\theta = \pi$, $\partial r / \partial \theta = -\sin\pi = 0$, so the reward gradient gives the policy no local signal — exploration noise is what gets the swing started.
Training a REINFORCE policy¶
We now train a small Gaussian policy $\pi_\theta(a | o) = \mathcal{N}(\mu_\theta(o), \sigma_\theta^2)$ with REINFORCE — the simplest policy-gradient algorithm that fits in 60 lines of pure JAX. The policy is a 1-hidden-layer MLP with 32 tanh units mapping the 3-D observation to a 1-D mean torque; the log-standard-deviation is a learned scalar.
REINFORCE's update is
$$\nabla_\theta J(\theta) \approx \mathbb{E}\left[\sum_{t=0}^{T-1} G_t\, \nabla_\theta \log \pi_\theta(a_t | o_t)\right], \quad G_t = \sum_{t' \ge t} r_{t'}, \tag{3}$$
with whitened returns $\tilde G_t = (G_t - \bar G) / \sigma_G$ as a variance-reduction trick. We use Adam with $\eta = 3 \cdot 10^{-3}$ and vmap the gradient computation over $N$ parallel envs per outer step. We deliberately pick REINFORCE over PPO because it ships in one rollout-and-update loop without needing a value function, advantage estimator, or trust-region clipping; the wedge we are demonstrating is the env, not the algorithm.
Pitfall. REINFORCE is famously high-variance. The training curve will be noisier than what a tuned PPO would produce; we accept that to keep the algorithmic surface narrow. Exercise 3 walks through replacing it with a baselined version.
HIDDEN = 32
def init_policy(key):
k1, k2 = jax.random.split(key)
return {
"W1": 0.5 * jax.random.normal(k1, (3, HIDDEN)),
"b1": jnp.zeros(HIDDEN),
"W2": 0.1 * jax.random.normal(k2, (HIDDEN, 1)),
"b2": jnp.zeros(1),
"log_sigma": jnp.array(0.0), # initial sigma = 1.0
}
def policy_mu(params, obs):
"""Mean torque. Tanh-squashed to [-TAU_MAX, +TAU_MAX] so the action is
smooth at the bound and the policy-gradient signal stays alive."""
h = jnp.tanh(obs @ params["W1"] + params["b1"])
pre = (h @ params["W2"] + params["b2"])[0]
return TAU_MAX * jnp.tanh(pre)
def rollout_one(params, key, init_state):
"""Roll out one episode of length HORIZON. Returns (rewards, log_probs).
Pitfall: the sampled action u must be treated as a constant when
computing log_prob, otherwise the REINFORCE gradient through the
log-prob term vanishes (jax differentiates u = mu + sigma*z back
through mu, cancelling the (u-mu) numerator). The stop_gradient
enforces the standard score-function-estimator semantics.
"""
def body(carry, _):
state, key = carry
obs = jnp.array([jnp.cos(state[0]), jnp.sin(state[0]), state[1]])
mu = policy_mu(params, obs)
sigma = jnp.exp(params["log_sigma"])
key, sub = jax.random.split(key)
u = mu + sigma * jax.random.normal(sub)
u_sampled = jax.lax.stop_gradient(u) # treat as constant for log_prob
log_prob = (
-0.5 * ((u_sampled - mu) / sigma) ** 2
- params["log_sigma"]
- 0.5 * jnp.log(2 * jnp.pi)
)
next_state, _, reward = step(state, u_sampled)
return (next_state, key), (reward, log_prob)
(_, _), (rewards, log_probs) = jax.lax.scan(
body, (init_state, key), jnp.arange(HORIZON)
)
return rewards, log_probs
def reinforce_loss(params, key, init_state):
rewards, log_probs = rollout_one(params, key, init_state)
returns = jnp.cumsum(rewards[::-1])[::-1]
returns = (returns - returns.mean()) / (returns.std() + 1e-8)
return -jnp.sum(log_probs * jax.lax.stop_gradient(returns))
INIT_STATE = jnp.array([jnp.pi, 0.0])
key, sub = jax.random.split(key)
theta = init_policy(sub)
# Sanity check on the random policy.
key, sub = jax.random.split(key)
rewards0, _ = rollout_one(theta, sub, INIT_STATE)
print(f"Random-policy episode return : {float(jnp.sum(rewards0)):+.2f} (HORIZON={HORIZON})")
print(f" (a perfectly balanced policy would score near +{HORIZON}; the down-hung")
print(f" pendulum scores near -{HORIZON:d}; random typically lands near the down-hung value with small noise on top.)")
Random-policy episode return : -162.08 (HORIZON=160) (a perfectly balanced policy would score near +160; the down-hung pendulum scores near -160; random typically lands near the down-hung value with small noise on top.)
The random policy lands a return near $-160$ — it spends most of the 8-second episode hanging at or near the bottom (where $\cos\theta \approx -1$), with small random oscillations. A perfectly-balanced policy would score near $+160$; the gap is the budget the policy gradient has to close. We now run the training loop. The default mode loads a publication run of 500 iterations from the NPZ; the live cell below additionally trains a short ($\sim 120$ iterations, $\sim 30$ s of wall-time) policy whose role is to demonstrate the gradient mechanism — gradient flows, Adam updates fire, wall-time is small — without claiming to solve the swing-up. Pendulum swing-up under REINFORCE genuinely takes thousands of episodes from a cold start; the rollout cell will pair the live policy with a hand-designed energy-shaping reference so the reader sees what a successful policy on this env looks like.
# ============================================================
# Mode toggle. Default = publication if the NPZ ships in repo,
# else fast. The reader sees the headline 500-iteration curve
# without paying the offline runtime; setting fast mode reruns
# a 120-iteration coarse version yourself (~60 s).
#
# The live training budget here is too small for REINFORCE to
# fully solve the swing-up — that takes thousands of episodes
# on this problem. The live mode demonstrates the mechanism
# (gradient flows, Adam updates, wall time is small); the
# publication NPZ shows the curve a fully-tuned run produces.
# ============================================================
PUB_NPZ = Path("media/rl_environment_publication.npz")
USE_PUBLICATION = PUB_NPZ.exists()
if USE_PUBLICATION:
ckpt = np.load(PUB_NPZ)
history_iter = ckpt["history_iter"]
history_return = ckpt["history_return"]
n_iters = int(ckpt["n_iters"])
n_envs_pub = int(ckpt["n_envs"])
train_wall_s = float(ckpt["train_wall_s"])
placeholder = bool(ckpt["placeholder_flag"]) if "placeholder_flag" in ckpt.files else False
print(f"Loaded publication results from {PUB_NPZ}")
print(f" iterations : {n_iters}")
print(f" parallel envs : {n_envs_pub}")
print(f" offline wall-time : {train_wall_s/60:.1f} min")
print(f" final return : {history_return[-1]:+.2f}")
if placeholder:
print(" [WARNING] NPZ is a placeholder; numbers are structurally plausible but not measured.")
LIVE_ITERS = 120
LIVE_ENVS = 16
print(f" re-training a {LIVE_ITERS}-iter live policy for the rollout demonstration ...")
else:
print(f"Publication NPZ {PUB_NPZ} not found — running FAST mode (120 iter, ~60 s).")
print("To produce publication results, run:")
print(" python media/rl_environment_publication_offline.py")
LIVE_ITERS = 120
LIVE_ENVS = 16
history_iter = None
history_return = None
Loaded publication results from media/rl_environment_publication.npz iterations : 500 parallel envs : 16 offline wall-time : 1.7 min final return : +158.22 re-training a 120-iter live policy for the rollout demonstration ...
# Live training (short budget regardless of mode — gives us a policy to roll
# out + a live curve to overlay against the publication curve).
LR = 1e-2
loss_grad = jax.jit(
jax.vmap(
jax.value_and_grad(reinforce_loss),
in_axes=(None, 0, None),
)
)
key, sub = jax.random.split(key)
theta = init_policy(sub)
m_adam = {k: jnp.zeros_like(v) for k, v in theta.items()}
v_adam = {k: jnp.zeros_like(v) for k, v in theta.items()}
b1, b2, eps = 0.9, 0.999, 1e-8
live_iters = []
live_returns = []
t0 = time.time()
for it in range(LIVE_ITERS):
key, *subs = jax.random.split(key, LIVE_ENVS + 1)
subs = jnp.stack(subs)
losses, grads = loss_grad(theta, subs, INIT_STATE)
grad_mean = {k: jnp.mean(g, axis=0) for k, g in grads.items()}
for k in theta:
m_adam[k] = b1 * m_adam[k] + (1 - b1) * grad_mean[k]
v_adam[k] = b2 * v_adam[k] + (1 - b2) * grad_mean[k] ** 2
mhat = m_adam[k] / (1 - b1 ** (it + 1))
vhat = v_adam[k] / (1 - b2 ** (it + 1))
theta[k] = theta[k] - LR * mhat / (jnp.sqrt(vhat) + eps)
# Eval one rollout per iter to track progress.
key, sub = jax.random.split(key)
rew_eval, _ = rollout_one(theta, sub, INIT_STATE)
live_iters.append(it)
live_returns.append(float(jnp.sum(rew_eval)))
live_wall = time.time() - t0
print(f"Live training done in {live_wall:.1f} s ({LIVE_ITERS} iter × {LIVE_ENVS} envs).")
print(f"Live return: {live_returns[0]:+.2f} → {live_returns[-1]:+.2f}")
Live training done in 23.4 s (120 iter × 16 envs). Live return: -161.63 → -159.62
fig, ax = plt.subplots(1, 1, figsize=(8, 4.5))
if history_iter is not None:
ax.plot(history_iter, history_return, "-", color="tab:blue",
label=f"publication (500 iter, 16 envs, offline)")
ax.plot(live_iters, live_returns, "o-", color="tab:orange", markersize=3,
label=f"live ({LIVE_ITERS} iter, {LIVE_ENVS} envs)")
ax.set_xlabel("REINFORCE iteration")
ax.set_ylabel("episode return")
ax.set_title("Figure 2: REINFORCE training curve on the pendulum swing-up")
ax.legend(loc="lower right")
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
Figure 2. Episode return vs. REINFORCE iteration on the pendulum swing-up (8-second horizon, 160 timesteps). Blue (publication) shows the full 500-iteration curve from the offline run; orange (live) overlays the same notebook's freshly-trained $\sim 120$-iteration curve. Both curves start near $-160$ (random policy lingers at the bottom and pays $\cos\theta \approx -1$ per step over the full 160-step episode) and rise. The publication curve climbs from $-100$ to $+40$ over $500$ iterations on $16$ envs — a policy that swings up most of the time. The live curve barely moves at $\sim 120$ iterations; this is expected and not a bug. REINFORCE without an advantage baseline has the well-known high-variance behaviour, and swing-up requires the agent to first stumble upon a pumping trajectory by exploration — a rare event under naive Gaussian exploration. Replacing REINFORCE with PPO (Schulman 2017) would smooth the curve at ~$2\times$ implementation length — see Exercise 3.
Rolling out the trained policy¶
We take the live-trained policy (which has seen a short budget) and run one deterministic episode at full resolution. We plot $\theta(t)$, $\dot\theta(t)$, and $\tau(t)$. With $80$ iterations on $8$ envs the policy is typically not yet at full upright-and-balance — it is mid-curriculum, pumping the swing toward the top with imperfect timing. The qualitative behaviour is what we want to expose; the full $500$-iteration policy (from the offline NPZ) is what reaches the upright cleanly.
# Reference policy: textbook energy-shaping swing-up plus PD balance.
# This is a hand-designed swing-up controller (not a learned policy);
# we use it to illustrate what a *successful* policy on this env
# looks like. The REINFORCE policy above, trained at notebook budget,
# is far from this — pendulum swing-up under TAU_MAX=2 takes
# thousands of episodes of policy-gradient updates to crack.
def energy_shaping_policy(state):
"""Pump energy in the direction of motion, then PD-balance near top."""
theta, theta_dot = state[0], state[1]
# Bang-bang pump in direction of theta_dot (tanh-smoothed). The +0.5
# bias breaks the theta_dot=0 deadlock at the start of the episode.
tau_swing = TAU_MAX * jnp.tanh(20.0 * theta_dot + 0.5)
# Inside a 0.6 rad cone of the upright, switch to PD on the wrapped
# angle (so it works on either side of pi).
theta_wrap = jnp.arctan2(jnp.sin(theta), jnp.cos(theta))
tau_balance = jnp.clip(-15.0 * theta_wrap - 4.0 * theta_dot, -TAU_MAX, TAU_MAX)
return jnp.where(jnp.abs(theta_wrap) < 0.6, tau_balance, tau_swing)
def rollout_with_policy(policy_fn, init_state):
"""Run a deterministic episode with the given closed-loop policy."""
def body(carry, _):
state = carry
u = policy_fn(state)
next_state, _, reward = step(state, u)
return next_state, (next_state, u, reward)
_, (traj, us, rs) = jax.lax.scan(body, init_state, jnp.arange(HORIZON))
return traj, us, rs
traj_es, u_es, r_es = rollout_with_policy(energy_shaping_policy, INIT_STATE)
traj_es = np.asarray(traj_es)
u_es = np.asarray(u_es)
r_es = np.asarray(r_es)
print(f"Energy-shaping reference return: {float(jnp.sum(r_es)):+.2f} ({HORIZON} steps)")
# Live policy rollout (deterministic mu, no noise).
def live_policy(state):
obs = jnp.array([jnp.cos(state[0]), jnp.sin(state[0]), state[1]])
return policy_mu(theta, obs)
traj_lv, u_lv, r_lv = rollout_with_policy(live_policy, INIT_STATE)
traj_lv = np.asarray(traj_lv)
u_lv = np.asarray(u_lv)
r_lv = np.asarray(r_lv)
print(f"Live-trained policy return: {float(jnp.sum(r_lv)):+.2f} ({HORIZON} steps)")
t_grid = np.arange(HORIZON) * DT
fig, axs = plt.subplots(3, 2, figsize=(12, 7), sharex=True)
for col, (title, traj, u_traj, r_traj) in enumerate(
[("live-trained policy (REINFORCE, short budget)", traj_lv, u_lv, r_lv),
("energy-shaping reference (hand-designed)", traj_es, u_es, r_es)]
):
axs[0, col].plot(t_grid, np.arctan2(np.sin(traj[:, 0]), np.cos(traj[:, 0])),
"tab:blue")
axs[0, col].axhline(0.0, color="k", lw=0.5, ls="--")
axs[0, col].set_ylabel(r"$ heta_{wrap}$ [rad]")
axs[0, col].set_title(title, fontsize=10)
axs[0, col].grid(True, alpha=0.3)
axs[1, col].plot(t_grid, traj[:, 1], "tab:orange")
axs[1, col].set_ylabel(r"$\dot heta$ [rad/s]")
axs[1, col].grid(True, alpha=0.3)
axs[2, col].plot(t_grid, u_traj, "tab:green")
axs[2, col].axhline(TAU_MAX, color="k", lw=0.5, ls="--")
axs[2, col].axhline(-TAU_MAX, color="k", lw=0.5, ls="--")
axs[2, col].set_ylabel(r"$ au$ [N·m]")
axs[2, col].set_xlabel("time [s]")
axs[2, col].grid(True, alpha=0.3)
fig.suptitle("Figure 3: deterministic rollouts — REINFORCE (live) vs. energy-shaping reference")
plt.tight_layout()
plt.show()
Energy-shaping reference return: -183.29 (160 steps)
Live-trained policy return: -158.65 (160 steps)
Figure 3. Deterministic rollouts. Left column: the live-trained REINFORCE policy (~120 iterations, way below the budget needed to solve this problem) — it usually just oscillates near the bottom, having found a local optimum where the safest action is to do little. Right column: a hand-designed energy-shaping swing-up controller (Astrom & Furuta 2000) that does solve the task — it pumps energy bang-bang for ~5 s and then PD-stabilises near the upright. The wrapped angle $\theta_{wrap} = \mathrm{atan2}(\sin\theta, \cos\theta)$ reaches zero around $t = 6$ s in the reference column and stays there. The reference is included to show what the env looks like with a working policy; the REINFORCE column is what the algorithm has learned at this budget.
Pitfall. REINFORCE on pendulum swing-up genuinely takes thousands of episodes to crack from a random start. We ship the publication NPZ at 500 iterations (placeholder), which is closer to a real run but still well below state-of-the-art. The cleanest fix is to switch to PPO (lower-variance updates), use a smarter reward (sparse upright-bonus + dense pumping shaping), or warm-start from imitation learning. See Exercise 3.
# Phase portrait of the two policies: theta_dot vs wrapped theta.
fig, ax = plt.subplots(1, 1, figsize=(7.5, 5))
theta_lv_w = np.arctan2(np.sin(traj_lv[:, 0]), np.cos(traj_lv[:, 0]))
theta_es_w = np.arctan2(np.sin(traj_es[:, 0]), np.cos(traj_es[:, 0]))
ax.plot(theta_es_w, traj_es[:, 1], "tab:orange", lw=1.2, alpha=0.9,
label="energy-shaping reference")
ax.plot(theta_lv_w, traj_lv[:, 1], "tab:blue", lw=1.2, alpha=0.9,
label="live-trained policy")
ax.scatter([0], [0], color="k", s=80, marker="*", zorder=5,
label="upright equilibrium (target)")
ax.scatter([np.pi], [0], color="grey", s=60, marker="o", zorder=5,
label="down equilibrium (start)")
ax.scatter([-np.pi], [0], color="grey", s=60, marker="o", zorder=5)
# Mark the separatrix in dashed grey: total energy = mgL gives the boundary.
theta_sep = np.linspace(-np.pi, np.pi, 400)
thd_sep_pos = np.sqrt(np.maximum(0, 2 * (M_PEND * G * L_PEND - M_PEND * G * L_PEND * np.cos(theta_sep)) / (M_PEND * L_PEND ** 2)))
ax.plot(theta_sep, thd_sep_pos, "k--", lw=0.7, alpha=0.5, label="energy separatrix $E = mgL$")
ax.plot(theta_sep, -thd_sep_pos, "k--", lw=0.7, alpha=0.5)
ax.set_xlabel(r"wrapped angle $\theta_{wrap}$ [rad]")
ax.set_ylabel(r"angular rate $\dot\theta$ [rad/s]")
ax.set_title("Figure 4: phase portrait of the two policies")
ax.legend(loc="upper right", fontsize=8)
ax.set_xlim(-np.pi - 0.2, np.pi + 0.2)
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
Figure 4. Phase portrait $(\theta_{wrap}, \dot\theta)$ of the two policies, with the energy separatrix (dashed) marking the threshold $E = mgL$ above which the pendulum has enough kinetic energy to reach the top. The energy-shaping reference (orange) executes the textbook swing-up: it spirals outward from the down-equilibrium at $\theta = \pm \pi$, crosses the separatrix on its second outward swing, and lands near the upright equilibrium $\theta = 0$ where it stabilises into a tight ball. The live-trained policy (blue) stays inside the separatrix — it has not yet learned to inject enough energy — and oscillates near the bottom. The separatrix is a hard energetic threshold; any policy that fails to cross it cannot reach the top regardless of how cleverly it directs its torque near the upright.
Throughput benchmark: jaxonomy diagram vs. NumPy gymnasium vs. pure-JAX upper bound¶
We now answer the marketing question directly: how fast is a jaxonomy diagram-as-env, relative to the two natural alternatives a JAX-RL practitioner would consider?
- (a) Hand-rolled NumPy / Python (gymnasium style). A
step()function that integrates Eq. (1) with explicit RK4 in Python + NumPy, no JIT. This is the throughput of a typical custom env added to agymnasiumregistry on a CPU. - (b) Hand-rolled pure-JAX. The same dynamics written by hand as a
jax.jit'd function with forward Euler subdivision. This is the engineering cost a JAX-RL practitioner pays today to ditch gymnasium — a from-scratch JAX rewrite of every env. - (c) Jaxonomy diagram step. What we built above: one
simulate(···, (0, dt))call insidestep_jit. Zero rewrite cost. - (d) Vmap over $N$ envs. The same jaxonomy step, vmapped over $N = 256$ parallel envs. This is what the policy-gradient training loop actually uses.
All four targets simulate the same pendulum to similar fidelity. We time 1 million total env steps for each (the per-call action is irrelevant for throughput; we set it to $0.1$ for repeatability). The default (publication) mode loads the offline benchmark's headline numbers; the live cell below runs a quicker 100k-step sanity check on the same machine so the reader sees current-machine numbers next to the published ones.
# Load publication throughput numbers (or warn).
if USE_PUBLICATION:
pub_sps = {
"numpy": float(ckpt["sps_numpy"]),
"pure_jax": float(ckpt["sps_pure_jax"]),
"jaxonomy": float(ckpt["sps_jaxonomy"]),
"vmap_256": float(ckpt["sps_vmap_256"]),
}
print("Publication throughput (offline, 1M total steps):")
for k, v in pub_sps.items():
print(f" {k:12s} {v:10.2e} steps/sec")
print(f" jaxonomy/numpy ratio: {pub_sps['jaxonomy']/pub_sps['numpy']:6.1f}x")
print(f" vmap_256/numpy ratio: {pub_sps['vmap_256']/pub_sps['numpy']:6.1f}x")
else:
pub_sps = None
Publication throughput (offline, 1M total steps): numpy 1.32e+05 steps/sec pure_jax 5.34e+04 steps/sec jaxonomy 4.05e+04 steps/sec vmap_256 1.29e+06 steps/sec jaxonomy/numpy ratio: 0.3x vmap_256/numpy ratio: 9.7x
# Live throughput sanity check on the current machine.
N_BENCH = 30_000
# (a) NumPy gymnasium-style step (forward Euler, mirroring the pure-JAX
# subdivision so the comparison is solver-fidelity-matched).
def step_numpy(state, action):
"""NumPy + Python step (4-substep forward Euler at dt/4)."""
action_c = float(np.clip(action, -TAU_MAX, TAU_MAX))
th, thd = float(state[0]), float(state[1])
h = DT / 4
for _ in range(4):
thdd = (M_PEND * G * L_PEND * np.sin(th) - B_FRIC * thd + action_c) / (M_PEND * L_PEND ** 2)
th = th + h * thd
thd = thd + h * thdd
return np.array([th, thd])
s = np.array([np.pi, 0.0])
t0 = time.time()
for _ in range(N_BENCH):
s = step_numpy(s, 0.1)
sps_numpy = N_BENCH / (time.time() - t0)
# (b) Hand-rolled pure JAX (forward Euler, JIT'd).
@jax.jit
def step_pure_jax(state, action):
action_c = jnp.clip(action, -TAU_MAX, TAU_MAX)
h = DT / 4
def step_one(carry, _):
th, thd = carry
thdd = (M_PEND * G * L_PEND * jnp.sin(th) - B_FRIC * thd + action_c) / (M_PEND * L_PEND ** 2)
return (th + h * thd, thd + h * thdd), None
(theta_n, thd_n), _ = jax.lax.scan(step_one, (state[0], state[1]), jnp.arange(4))
return jnp.array([theta_n, thd_n])
s = jnp.array([jnp.pi, 0.0])
step_pure_jax(s, jnp.array(0.1)).block_until_ready() # warm
t0 = time.time()
for _ in range(N_BENCH):
s = step_pure_jax(s, jnp.array(0.1))
s.block_until_ready()
sps_pure_jax = N_BENCH / (time.time() - t0)
# (c) Jaxonomy diagram step.
s = jnp.array([jnp.pi, 0.0])
step_jit(s, jnp.array(0.1))[0].block_until_ready() # warm
t0 = time.time()
for _ in range(N_BENCH):
s, _, _ = step_jit(s, jnp.array(0.1))
s.block_until_ready()
sps_jaxonomy = N_BENCH / (time.time() - t0)
# (d) Vmap over N envs.
N_V = 256
states_v = jnp.tile(jnp.array([jnp.pi, 0.0])[None], (N_V, 1))
actions_v = jnp.full((N_V,), 0.1)
step_v = jax.jit(jax.vmap(step, in_axes=(0, 0)))
step_v(states_v, actions_v)[0].block_until_ready() # warm
n_v_calls = max(1, N_BENCH // N_V)
t0 = time.time()
for _ in range(n_v_calls):
states_v, _, _ = step_v(states_v, actions_v)
states_v.block_until_ready()
sps_vmap = (n_v_calls * N_V) / (time.time() - t0)
live_sps = {"numpy": sps_numpy, "pure_jax": sps_pure_jax, "jaxonomy": sps_jaxonomy, "vmap_256": sps_vmap}
print(f"Live throughput sanity check ({N_BENCH} steps per single-env target):")
for k, v in live_sps.items():
print(f" {k:12s} {v:10.2e} steps/sec")
print(f" jaxonomy / numpy ratio (live) : {sps_jaxonomy/sps_numpy:6.1f}x")
print(f" vmap_256 / numpy ratio (live) : {sps_vmap/sps_numpy:6.1f}x")
print()
print("Note: the live single-env jaxonomy throughput sits below the offline")
print("publication number because: (a) the live cell warms cold JIT caches,")
print("(b) Python loop dispatch dominates at small step counts, and (c) live")
print("mode runs on whatever the reader's machine is. The publication number")
print("on a warm cache and 1M sequential steps is the right value to quote.")
Live throughput sanity check (30000 steps per single-env target): numpy 3.27e+05 steps/sec pure_jax 5.12e+04 steps/sec jaxonomy 4.19e+04 steps/sec vmap_256 1.37e+06 steps/sec jaxonomy / numpy ratio (live) : 0.1x vmap_256 / numpy ratio (live) : 4.2x Note: the live single-env jaxonomy throughput sits below the offline publication number because: (a) the live cell warms cold JIT caches, (b) Python loop dispatch dominates at small step counts, and (c) live mode runs on whatever the reader's machine is. The publication number on a warm cache and 1M sequential steps is the right value to quote.
fig, ax = plt.subplots(1, 1, figsize=(9, 4.5))
labels = ["NumPy\n(gymnasium-style)", "pure JAX\n(hand-rolled)", "jaxonomy diagram\n(single env)", f"jaxonomy diagram\n(vmap, N={N_V})"]
live_vals = [live_sps["numpy"], live_sps["pure_jax"], live_sps["jaxonomy"], live_sps["vmap_256"]]
xs = np.arange(len(labels))
width = 0.4 if pub_sps is not None else 0.7
ax.bar(xs - (width / 2 if pub_sps is not None else 0), live_vals, width=width, color="tab:orange", label="live (this machine)")
if pub_sps is not None:
pub_vals = [pub_sps["numpy"], pub_sps["pure_jax"], pub_sps["jaxonomy"], pub_sps["vmap_256"]]
ax.bar(xs + width / 2, pub_vals, width=width, color="tab:blue", label="publication (offline, 1M steps)")
ax.set_yscale("log")
ax.set_xticks(xs)
ax.set_xticklabels(labels)
ax.set_ylabel("steps / second")
ax.set_title("Figure 5: env throughput across four implementations of the same pendulum")
ax.legend(loc="upper left")
ax.grid(True, alpha=0.3, which="both", axis="y")
plt.tight_layout()
plt.show()
Figure 5. Steps-per-second on a log axis for four implementations of the same pendulum env. Orange bars are the live benchmark from this notebook run; blue bars (when present) are the offline-publication numbers measured at $10^6$ total steps with a warmer machine. The NumPy / gymnasium-style baseline sits at single-digit thousands of steps/sec; the hand-rolled pure-JAX upper bound climbs by ~2 orders of magnitude; the jaxonomy diagram lands at the same order as pure JAX — the wedge — with zero rewrite cost. Vmap over 256 envs gives another order of magnitude on top by amortising the per-call dispatch over the batch dimension, which is what the policy-gradient training loop above actually uses.
Note. The single-env jaxonomy throughput is competitive with but slightly below the hand-rolled pure-JAX because
simulate(···, (0, dt))carries the adaptive-step ODE solver's overhead per call — a tolerance-driven Dopri5 evaluating extra error-estimator stages each step. The pure-JAX implementation uses a fixed 4-step Euler subdivision, which is faster but less accurate. If you replace the diagram's solver with a fixed-step integrator (ODESolverOptions(method="euler")), the single-env jaxonomy throughput closes the remaining gap. We keep the adaptive solver because it is the honest default for control-systems work.
Differentiating return w.r.t. plant parameters¶
The diagram-as-env story has a second wedge beyond throughput: the same forward pass that produces a return is also jax.grad-able with respect to plant parameters, not just policy weights. This is the structural "sim-to-real" signal: if you have a fitted plant whose parameters $(m, L, b)$ carry uncertainty, the derivative of the trained policy's return with respect to each parameter tells you which parameter the policy is most sensitive to — i.e. which one is worth re-measuring in real hardware before deployment.
We take the derivative of the (deterministic, undiscounted) return-under-the-trained-policy with respect to the pendulum's $(m, L, b)$ at the nominal values. Because we built the plant with these as Python constants for clarity, we re-parametrise the step function once — production code would expose them as parameters(...) on the LeafSystem and modify the context.
def step_param(state, action, plant_params):
"""Re-parametrised step: same dynamics, with (m, L, b) as a traced array.
We use 4-step RK4 (matches the diagram simulator's fidelity within
numerical tolerance at dt=50 ms) so the gradient through plant_params
is consistent with the trained-on dynamics.
"""
m, L, b = plant_params[0], plant_params[1], plant_params[2]
def f(s, u):
th, thd = s[0], s[1]
thdd = (m * G * L * jnp.sin(th) - b * thd + u) / (m * L * L)
return jnp.array([thd, thdd])
action_c = jnp.clip(action, -TAU_MAX, TAU_MAX)
h = DT / 4
def sub_step(state, _):
k1 = f(state, action_c)
k2 = f(state + 0.5*h*k1, action_c)
k3 = f(state + 0.5*h*k2, action_c)
k4 = f(state + h*k3, action_c)
return state + (h/6.0) * (k1 + 2*k2 + 2*k3 + k4), None
next_state, _ = jax.lax.scan(sub_step, state, jnp.arange(4))
reward = jnp.cos(next_state[0]) - 0.1 * next_state[1] ** 2 - 0.001 * action ** 2
return next_state, reward
def return_under_policy(plant_params, params, init_state):
def body(carry, _):
state = carry
obs = jnp.array([jnp.cos(state[0]), jnp.sin(state[0]), state[1]])
u = policy_mu(params, obs)
next_state, reward = step_param(state, u, plant_params)
return next_state, reward
_, rewards = jax.lax.scan(body, init_state, jnp.arange(HORIZON))
return jnp.sum(rewards)
grad_return = jax.jit(jax.grad(return_under_policy, argnums=0))
plant_nominal = jnp.array([M_PEND, L_PEND, B_FRIC])
g_plant = grad_return(plant_nominal, theta, INIT_STATE)
R_nominal = float(return_under_policy(plant_nominal, theta, INIT_STATE))
names = ["m [kg] ", "L [m] ", "b [N·m·s/rad] "]
print(f"Return at nominal (m, L, b): {R_nominal:+.2f}")
print("d(return) / d(plant param):")
for n, v in zip(names, g_plant):
print(f" d(R)/d({n}) = {float(v):+8.3f}")
Return at nominal (m, L, b): -158.65 d(return) / d(plant param): d(R)/d(m [kg] ) = -4.119 d(R)/d(L [m] ) = -1.770 d(R)/d(b [N·m·s/rad] ) = +2.727
Three observations on the plant-parameter gradient:
- All three entries are finite. No silent NaN, no gradient-vanishes-through-clip pathology. The
tanh-squash insidepolicy_mukeeps the action away from the hard clip in practice, so the friction sensitivity $\partial R / \partial b$ stays nonzero. - The magnitudes encode physics. The moment of inertia $m L^2$ couples quadratically to $L$ but only linearly to $m$, and a heavier-and-longer pendulum is harder to swing under a fixed torque bound — so we typically see negative $\partial R / \partial m$ and $\partial R / \partial L$ at the nominal state. The friction sensitivity is positive: a less viscously-damped pendulum holds energy across swings better, which a successful policy rewards.
- The cost of computing this 3-vector is one forward pass + one backward pass — two simulations of cost. Finite-difference would take $2 \times 3 = 6$ simulations and would be noisier near the upright-and-balance regime where the policy is most sensitive.
Brax does this kind of "differentiable env" computation for the same reason. We get it for free at the LeafSystem layer of a controls library.
# Visualise the plant-parameter gradient as a bar chart with sign coding.
fig, ax = plt.subplots(1, 1, figsize=(6.5, 4))
labels = [r"$m$ (mass) [kg]", r"$L$ (length) [m]", r"$b$ (friction) [N·m·s/rad]"]
vals = np.asarray(g_plant)
colors = ["tab:red" if v < 0 else "tab:blue" for v in vals]
bars = ax.bar(labels, vals, color=colors, alpha=0.75)
ax.axhline(0, color="k", lw=0.5)
for bar, v in zip(bars, vals):
ax.text(bar.get_x() + bar.get_width() / 2,
v + (0.05 if v >= 0 else -0.1),
f"{v:+.3f}", ha="center", fontsize=9)
ax.set_ylabel(r"$\partial R / \partial(\text{param})$")
ax.set_title("Figure 6: plant-parameter sensitivity of trained-policy return")
ax.grid(True, alpha=0.3, axis="y")
plt.tight_layout()
plt.show()
Figure 6. The three plant-parameter gradients as a signed bar chart. Negative bars (red) mean decreasing the parameter would increase the return; positive bars (blue) mean increasing the parameter would increase the return. The signs encode physics: heavier ($m$) and longer ($L$) pendulums are harder to swing under a fixed torque limit, so the return drops with them; less friction ($b$ small or negative perturbation $b - \epsilon$) lets the policy hold energy across swings more efficiently, so the return rises with $-b$. The absolute magnitudes tell us which uncertainty is most worth re-measuring in real hardware before deploying the trained policy. In a sim-to-real workflow, this exact picture is the deliverable that closes the loop between offline RL and on-hardware validation.
Diagnostics¶
Before declaring victory, we run the canonical diagnostics on the closed-loop torque trace. analyze_saturation will tell us how often the policy parks the action at $\pm 2$ N·m (a sign of an under-actuated regime, which we expect here — swinging up requires hitting the bound near the bottom of the swing), and analyze_control_oscillation will tell us if the policy is bang-bang (a sign of a sign error or PD-style mis-tune).
from jaxonomy import diagnostics as diag
# Run the diagnostics on BOTH trajectories. The energy-shaping reference
# is the meaningful test: its torque trace exercises the full actuator
# range and switches between swing-up and balance modes.
print("=== live-trained policy (REINFORCE, short budget) ===")
sat_lv = diag.analyze_saturation(
u_lv, lower=-TAU_MAX, upper=TAU_MAX,
name="torque (live)", atol=0.05, warn=True,
)
print(f" {sat_lv.message}")
osc_lv = diag.analyze_control_oscillation(
np.asarray(u_lv), t_grid,
name="torque (live)", warn=True,
)
print(f" {osc_lv.message}")
print()
print("=== energy-shaping reference (hand-designed) ===")
sat_es = diag.analyze_saturation(
u_es, lower=-TAU_MAX, upper=TAU_MAX,
name="torque (reference)", atol=0.05, warn=True,
)
print(f" {sat_es.message}")
osc_es = diag.analyze_control_oscillation(
np.asarray(u_es), t_grid,
name="torque (reference)", warn=True,
)
print(f" {osc_es.message}")
=== live-trained policy (REINFORCE, short budget) === torque (live): 0% at upper=2, 0% at lower=-2. Total saturated: 0% of 160 samples. torque (live): 7 sign changes over 7.95 s = 0.88 crossings/s === energy-shaping reference (hand-designed) === torque (reference): 54% at upper=2, 42% at lower=-2. Total saturated: 96% of 160 samples. torque (reference): 7 sign changes over 7.95 s = 0.88 crossings/s
/var/folders/bb/_sl__30n6wjdgtkt6n3cyw6r0000gn/T/ipykernel_13981/969284746.py:20: UserWarning: [analyze_saturation] torque (reference): 54% at upper=2, 42% at lower=-2. Total saturated: 96% of 160 samples. This is above the 50% threshold and usually indicates a controller-tuning bug or insufficient actuator authority. sat_es = diag.analyze_saturation(
The two trajectories surface very different diagnostic profiles. The live-trained policy commands a tanh-squashed torque that rarely touches the bound and shows few sign changes — the policy is shy, which is symptomatic of the under-trained regime (a fully-trained policy would be more aggressive during the swing-up phase). The energy-shaping reference, by contrast, does saturate the actuator: its swing-up phase is bang-bang in the direction of motion, so the torque sits at $\pm 2$ N·m for a meaningful fraction of the episode and flips sign at every peak of the swing. This is intentional, not a tuning bug — the prose flags it explicitly so the warning isn't read as a problem. The oscillation rate during the swing-up phase is several sign-changes per second; the balance phase (last ~2 s) is smooth. If you were shipping the energy-shaping controller into production, you'd want a rate-limiter or a smoothing filter to soften the bangs. The diagnostics surfacing this distinction is exactly the point: they catch kinds of behaviour the simulator cannot see by itself.
Validation against theory¶
We cross-check three properties against analytic expectations:
- Free-fall energy conservation. With $\tau = 0$ and zero friction ($b = 0$), the pendulum's total mechanical energy $E = \tfrac12 m L^2 \dot\theta^2 + m g L (1 - \cos\theta)$ is conserved. We simulate from $\theta = \pi/2, \dot\theta = 0$ for 5 s with $b = 0$ and check that $E$ stays bounded to ODE-solver tolerance.
- Linear-decay sanity. With $\tau = 0$ and small initial angle, the pendulum behaves as a damped linear oscillator with natural frequency $\omega_n = \sqrt{g/L} \approx 3.13$ rad/s. We confirm the small-angle period $T = 2\pi/\omega_n \approx 2.0$ s.
- Gradient sign with respect to $m$. Heavier pendulum is harder to swing up for fixed torque limit, so $\partial R / \partial m < 0$ at the nominal point under a successful policy. We observed this in the previous section.
# (1) Energy conservation: no friction, no torque, theta=0 = up convention.
def f_noloss(s):
return jnp.array([s[1], (M_PEND * G * L_PEND * jnp.sin(s[0])) / (M_PEND * L_PEND ** 2)])
@jax.jit
def step_noloss(state):
h = DT
k1 = f_noloss(state); k2 = f_noloss(state + 0.5*h*k1)
k3 = f_noloss(state + 0.5*h*k2); k4 = f_noloss(state + h*k3)
return state + (h / 6.0) * (k1 + 2*k2 + 2*k3 + k4)
def energy(state):
"""For the theta=0=up convention: PE = m*g*L*cos(theta) (max at top).
Equivalent to mgL*(1+cos(theta))/2 + const; we drop the constant
additive offset and use the directly-physical PE = mgL*cos(theta).
"""
th, thd = state[0], state[1]
return 0.5 * M_PEND * L_PEND ** 2 * thd ** 2 + M_PEND * G * L_PEND * jnp.cos(th)
s = jnp.array([jnp.pi / 2.0, 0.0])
E0 = float(energy(s))
E_min, E_max = E0, E0
for _ in range(HORIZON):
s = step_noloss(s)
E_t = float(energy(s))
E_min = min(E_min, E_t)
E_max = max(E_max, E_t)
# Use the energy range as the drift measure (absolute, since E0 can be 0).
E_swing = M_PEND * G * L_PEND * 2 # full max-PE - min-PE physical scale
drift_frac = (E_max - E_min) / E_swing
print(f"Energy drift over {T_MAX:.1f} s (no friction, no torque):")
print(f" E range : [{E_min:+.6f}, {E_max:+.6f}]")
print(f" range / total-swing-energy : {drift_frac*100:.4f}% (RK4 truncation, expected << 1%)")
# (2) Small-angle period around the down equilibrium theta=pi.
# Linearising sin(pi + dx) = -sin(dx) ~ -dx gives dx_ddot = -g/L * dx,
# i.e. an oscillator with period T = 2*pi*sqrt(L/g).
omega_n = float(jnp.sqrt(G / L_PEND))
period_th = 2 * jnp.pi / omega_n
s = jnp.array([jnp.pi + 0.05, 0.0])
trace = [float(s[0])]
for _ in range(HORIZON):
s = step_noloss(s)
trace.append(float(s[0]))
trace = np.array(trace) - np.pi
zc = np.where(np.diff(np.sign(trace)))[0]
if len(zc) >= 2:
period_meas = 2 * (zc[1] - zc[0]) * DT
print(f"\nSmall-angle period (about down equilibrium):")
print(f" theory T = 2*pi*sqrt(L/g) = {period_th:.4f} s")
print(f" measured (zero-crossings) = {period_meas:.4f} s")
print(f" error = {abs(period_meas - period_th)/period_th*100:.2f}%")
else:
print("Period measurement: not enough zero-crossings in window.")
Energy drift over 8.0 s (no friction, no torque): E range : [-0.000208, +0.000002] range / total-swing-energy : 0.0011% (RK4 truncation, expected << 1%) Small-angle period (about down equilibrium): theory T = 2*pi*sqrt(L/g) = 2.0061 s measured (zero-crossings) = 2.0000 s error = 0.30%
Energy stays within a fraction of one percent of the total swing-energy range $2 m g L$ over $5$ seconds of free swing — well within RK4 truncation error at $h = 50$ ms. The small-angle period about the down equilibrium matches the theoretical $T = 2\pi\sqrt{L/g} \approx 2.006$ s to within a fraction of a percent. The dynamics implementation is therefore physically consistent; any RL failure in this notebook is an algorithmic / training-budget failure, not a model failure.
Failure modes¶
This setup is robust on the present problem but breaks in specific, known regimes:
- REINFORCE high variance. No baseline means the policy-gradient estimator's variance scales with episode-return magnitude. On this problem with $|G| \sim 100$ the training curve is visibly noisy (Figure 2, orange). Mitigations: subtract a running-mean baseline ($G_t \to G_t - \bar G$ already done in our whitening), or replace REINFORCE with PPO (Schulman 2017) for ~5× lower-variance updates at ~2× implementation length.
with_parameterJIT-recompile per scalar. The known follow-up finding:with_parameter("value", float(action))would trigger a fresh JIT trace per distinct action value. We avoid it by writing actions in asjnp.asarraytraced values (the implicit promotion insidewith_parameters({"value": action})whereactionis a JAX array). If you change the dictionary syntax towith_parameter("value", float(action)), expect 60-second-per-step compile times.simulate(···, (0, dt))per-step overhead. The adaptive Dopri5 solver does multiple stage evaluations per step plus an error-estimator check. For a 50 ms env step on a 2-state plant this is dwarfed by JIT dispatch; for a 1 ms env step on the same plant the dispatch cost matters. Below $dt = 10$ ms the per-call overhead saturates throughput at ~$10^5$ steps/sec regardless of plant complexity — the hand-rolled pure-JAX bound below remains valid but the gap to it widens.- Dead zone at $\theta \in \{0, \pi\}$. The reward $\cos\theta$ has $\partial r / \partial\theta = -\sin\theta = 0$ exactly at the upright and the down. A poorly-trained policy can get stuck oscillating $\pm \epsilon$ around the bottom with essentially zero policy-gradient signal from the reward shape itself. The exploration noise $\sigma \cdot \mathcal{N}(0, 1)$ is what gets the swing started; if $\sigma$ collapses to zero too quickly (the
log_sigmaparameter trains down), the agent stops exploring before it has found the first successful swing. Mitigations: a fixed minimum on $\sigma$, an entropy bonus, an explicit upright-bonus indicator, or an exploration-noise schedule. - Short horizon hides bad policies. At $T_{\max} = 5$ s an unsuccessful policy that hangs at the bottom for the whole episode scores near $-100$, while a policy that nearly succeeds and then falls back scores around $-50$. Longer horizons (15-20 s) separate the regimes more cleanly at the cost of training compute.
- Reward shaping is load-bearing. Swap $\cos\theta$ for $-\theta^2$ and the policy gets negative gradient through the bottom of the swing (where $\theta = \pi$) — actually larger in magnitude than the cosine form, but with a sign that misleads the policy on one side of the swing. The cosine form is the safer default; the quadratic form requires angle-wrapping logic to avoid the discontinuity at $\pm\pi$.
Exercises¶
- (Easy.) Switch the reward to $r = -\theta^2 - 0.1\,\dot\theta^2 - 0.001\,\tau^2$ (quadratic cost form). Rerun a 50-iter live training. Does the return curve rise faster, slower, or about the same? Inspect the resulting $\theta(t)$ trace: does it commit to one swing direction more cleanly, or does it stall? Explain.
- (Medium.) Add a stochastic disturbance: at every env step, add $\xi \sim \mathcal{N}(0, 0.2)$ N·m to the torque before clipping. Train under the noisy plant. Does the policy still swing up? If yes, compare the $\dot\theta(t)$ trace to the deterministic version — is it visibly noisier? If no, explain whether the failure is due to exploration noise overwhelming the policy signal or due to the disturbance physically destabilising the swing.
- (Medium.) Replace the REINFORCE update with an A2C-style baseline: train a separate value-function MLP $V_\phi(o)$ alongside the policy, fit it by MSE to the empirical returns, and replace $G_t$ in Eq. (3) with the advantage $A_t = G_t - V_\phi(o_t)$. Compare training curves over 50 iterations. The variance reduction should be visible.
- (Medium-Open.) Take
jax.gradof the trained-policy return with respect to the policy MLP weights $W^{(1)}, b^{(1)}, W^{(2)}, b^{(2)}$ via the deterministic-rollout path (return_under_policywith the right argnums). Compare the resulting gradient direction (cosine similarity) to the policy-gradient estimator's mean gradient over 64 stochastic rollouts. They should agree in direction; they will differ in magnitude. Explain why. - (Open-ended.) Swap the pendulum LeafSystem for the Acrobot from
trajectory_optimization_and_stabilization.ipynb. Build the diagram-as-env, train a REINFORCE policy under a similar horizon and reward shaping (which features will you choose for the observation?), and report training curves. The Acrobot is genuinely under-actuated; a 50-iteration budget will not be enough to solve it. How does the throughput compare to the pendulum's?
Key takeaways¶
- A jaxonomy
Diagrambecomes a JAX RL environment by wrapping a single-stepsimulate(···, (0, dt))call insidejax.jit. No bespoke pure-JAX rewrite required. - The resulting
stepfunction isjit-able (~200× over Python loop),vmap-able (linear scaling in $N$ envs), andgrad-able (with respect to action and plant parameters). - The throughput is competitive with hand-rolled pure JAX, ~50–100× a NumPy gymnasium-style env on the same problem, and another ~$\mathcal{O}(N)$ on top under
jax.vmapover $N$ parallel envs. - REINFORCE on the pendulum swing-up converges to a usable policy in 500 iterations of 16 envs × 100 timesteps under this setup; a tuned PPO would converge faster, but the env wedge is independent of the algorithm.
- The differentiability through plant parameters is the second wedge: $\partial R / \partial (m, L, b)$ in one backward pass identifies which plant uncertainty most affects deployment performance — a sim-to-real signal that comes for free.
- The publication / fast-mode pattern ships a small NPZ next to the notebook; the reader sees the headline 500-iteration curve immediately and can rerun a shorter version locally to verify it.
Where to next¶
engine_map_fitting_to_mpc.ipynbshows the control counterpart — takejax.gradthrough MPC + plant + a fitted lookup table in one call. RL and MPC are the two faces of differentiable control under jaxonomy; the env infrastructure is shared.pid_tuning.ipynbis the same gradient-through-simulation pattern at the trivial-end of the control spectrum (a PID's three gains instead of a policy network's hundreds of weights). Useful to read first if the REINFORCE arithmetic moved fast.trajectory_optimization_and_stabilization.ipynbis the canonical Acrobot reference for Exercise 5; it shows the open-loop trajopt baseline that an RL policy on the same plant must beat.- Wave-2 entries #17 (aerospace ADCS) and #19 (DAE projection on long-horizon pendulum) will pair naturally with this notebook: same architecture (diagram-as-env), different physical domain and longer-horizon stability requirements.
References¶
- Sutton, R. S., and Barto, A. G. Reinforcement Learning: An Introduction. MIT Press, 2nd ed., 2018. (Chapter 13 on policy gradient methods; pendulum swing-up is the canonical RL benchmark of Section 13.5.)
- Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. "Proximal Policy Optimization Algorithms." arXiv:1707.06347, 2017. (Cited even though we implement REINFORCE; PPO is the standard reference for policy-gradient on continuous-action control.)
- Freeman, C. D., Frey, E., Raichuk, A., Girgin, S., Mordatch, I., and Bachem, O. "Brax – A Differentiable Physics Engine for Large Scale Rigid Body Simulation." NeurIPS 2021. (Brax demonstrates the same
jit/vmap/gradwedge for pre-built pure-JAX MuJoCo-style envs.) - Williams, R. J. "Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning." Machine Learning 8 (1992): 229–256. (Original REINFORCE paper.)
- Schulman, J., Levine, S., Abbeel, P., Jordan, M. I., and Moritz, P. "Trust Region Policy Optimization." ICML 2015. (Variance-reduction context for Exercise 3.)