balance_robot

PPO policies for a two-wheeled balancing robot (wheeled inverted pendulum), trained in MuJoCo Warp via mjlab with rsl_rl and cross-checked against a PyBullet oracle.

Each policy is an ONNX file laid out as <run>/model_<iter>.onnx. The run name is the training recipe; results for each are in the source repo's TRAINING_LOG.md. Older entries are raw rsl_rl .pt checkpoints (below).

Several observation interfaces live in this repo. The sk_* runs are the production interface: 40 inputs, ten channels over four frames. The abl_* runs are interface-ablation artifacts, and they differ from each other as well as from sk_*: abl_combo is 10 inputs wide, while abl_nolpf_jerk1 keeps all 40 and changes what one channel means. Read the width and the filter constants from each file's metadata rather than assuming; the section on ablation policies below says why that matters more than usual.

Running an .onnx policy

import json, numpy as np, onnx, onnxruntime as ort
m = onnx.load("sk_ident_r128_s0/model_4500.onnx")
meta = {p.key: p.value for p in m.metadata_props}
names, scales = json.loads(meta["obs_names"]), np.array(json.loads(meta["obs_scales"]))
sess = ort.InferenceSession(m.SerializeToString())
obs = raw_obs / scales                       # raw_obs in the order of `names`
action = sess.run(["action"], {"obs": obs[None].astype(np.float32)})[0][0]
U, u_y = action * json.loads(meta["action_scale_volts"])   # volts

The graph is the deterministic actor: tanh MLP, linear head, output clipped to [-1, 1], batch dimension free. Input obs is [N, obs_dim] — take the channel order and the count from obs_names, and obs_history for how many frames are stacked (oldest first). One frame is pitch, pos_err, gyro, odom_vel, gyro_z, cmd_vel, cmd_yaw_rate, yaw_err, prev_action[0], prev_action[1], each raw value divided by its scale (obs_scales).

Output action is [N, 2]; volts are action * action_scale_volts ([8, 4]) and the wheel mix is motor_left = -0.5*(U+u_y), motor_right = -0.5*(U-u_y). The metadata also carries the source checkpoint, run name, training config and git sha.

Ablation policies (abl_*)

These come from an interface-ablation ladder: which parts of the observation interface actually earn their keep? Each arm removes one piece entirely — never shrinks it — and trains on an otherwise identical plant, reward and PPO recipe, against sk_ident_r128_s0/model_4500 as the control.

abl_combo_r128_s0 removes four things at once and is the interesting one:

	control (sk_ident)	abl_combo
observation	40 (10 channels × 4 frames)	10 (1 frame)
odom_vel	50 ms low-pass	raw wheel-mean velocity × nominal radius
yaw_err source	fused heading estimator (τ 10 s, odometric PI, slip gate, settle clock)	bias-corrected gyro-Z integration only
prev-action noise	0.15 during training	0 (training-only; no deploy effect)
network	40 → 64 → 64 → 2	10 → 64 → 64 → 2

Everything else is identical: channel order, obs_scales, action_scale_volts, the wheel mix, 100 Hz policy rate.

Why it is worth having. Removing frame stacking alone costs 0.043 survival and 0.129 on the long-action-delay bin. Removing it together with the other three costs almost nothing — 0.970 against the control's 0.976, inside the ±0.011 seed-to-seed noise — with the calmest pitch trace in the ladder. The interface pieces are not independent, and single-arm ablations of a jointly-tuned interface mislead.

Read the interface metadata before feeding this policy anything. Because these arms differ in what the channels mean rather than only in width, each abl_* file carries extra metadata keys:

ablation_arm      combo
odom_lpf_tau_s    0.0
odom_vel_source   RAW wheel-mean velocity x nominal radius, NO low-pass
yaw_tau_s         0.0
yaw_err_source    bias-corrected gyro-Z integration ONLY -- no odometric
                  fusion, slip gate or settle clock

Feeding a combo policy a filtered odom_vel, or a fused heading, is a silent and severe failure rather than a small offset: the reverse mismatch — the 40-input control policy fed a raw velocity channel it never trained on — survives 0.000 of episodes.

abl_nolpf_jerk1

Same 40-input interface as the control, same fused heading, with one change: odom_vel is unfiltered. The 50 ms low-pass was removed and replaced during training by a penalty on the second difference of the action, |a_t - 2a_(t-1) + a_(t-2)|^2 at weight 1.0, testing whether the filter's denoising can live in the reward instead, where it costs nothing at deployment and needs no firmware constant. The penalty is training-only and does not appear in the interface.

It half worked. Pitch wobble at the 12 ms bench condition came out at 1.09 deg against the control's 1.58, so the reward term denoises better than the filter did. It cost survival: 0.965 against 0.975, and the 36-44 ms held-delay bin 0.891 against 0.923. Penalising rapid reversals suppresses dither, and it also suppresses the fast corrections needed when the control loop is long. A companion arm at weight 4.0 never learned to balance at all and was killed at iteration 700 with episode length 62.

Bench-tested 2026-09-11, and it is worse on the robot. It is more susceptible to disturbances and drives further forward before recovering, which is what the 36-44 ms bin predicted. Published as a record of the experiment. Use sk_ident_r128_s0 instead.

The reason looks structural rather than a matter of tuning. A penalty on the action cannot separate dither-driven jitter from the fast corrections a long control loop needs, because at the action they are the same signal. A companion arm at weight 4.0 suppressed both and never balanced; this one at 1.0 suppresses both mildly and recovers sluggishly. The low-pass works because it acts at the input, where the two are still separable by frequency.

Caveats. One seed per arm; the headline gap is inside single-seed noise, and the mjlab evaluator is itself non-deterministic (three runs of one checkpoint, same seed: 0.974 / 0.976 / 0.982). abl_combo's weakest axis is the 36–44 ms held-action-delay bin, 0.895 against the control's 0.936 — which is exactly the regime where hardware transfer is most fragile. These are research artifacts, not a recommended deployment.

Loading a raw .pt (older entries)

A checkpoint is a torch.save dict with actor_state_dict, critic_state_dict, optimizer_state_dict, iter and infos. Only the actor matters for deployment:

actor_state_dict:
  mlp.0.weight (64, obs_dim)   mlp.0.bias (64,)
  mlp.2.weight (64, 64)        mlp.2.bias (64,)
  mlp.4.weight (2, 64)         mlp.4.bias (2,)
  distribution.std_param (2,)     # exploration only; unused at deploy time

The odd indices are the activations of the nn.Sequential, so the layer indices run 0, 2, 4. obs_dim is 40 for the sk_* runs and 10 for abl_combo.

The checkpoint does not record the activation function, the input layout, or the filter constants the inputs assume — which is exactly why the published form is ONNX. All of it is fixed by the training config:

actor	MLP [64, 64], tanh activation
observation	layout sk; history and width per run (see obs_history)
action	2 values in [-1, 1], wheel volts as a fraction of the limit
control rate	100 Hz policy, 500 Hz physics

The deterministic action is the Gaussian mean — the last layer's raw output, no tanh — clipped to [-1, 1]. There is no observation normalization: a checkpoint carrying normalizer state does not belong here.