MuJoCo simulation of a pair of arms, wrapped as a Gymnasium environment. Physics steps on the GPU via MJX.
uv sync
Requires Python 3.14 or newer and a CUDA GPU. CUDA libraries install as dependencies.
uv run demo
Opens a viewer and plays tracking episodes at real time: each one sweeps a noisy reference from the ready pose to a fresh final target, drawn as a translucent green arm, and the arms follow it. Grey and yellow spheres mark the arm centre of mass at the episode start and now, and each episode prints its arrival error and how far the centre of mass moved. Drag to orbit, scroll to zoom.
The demo follows the reference exactly as commanded, with no learned policy and no regard for the centre of mass; evaluate --gui plays a trained policy on the same episodes.
uv sync
uv run pre-commit install
ruff formats and lints, pyrefly type checks, both as pre-commit hooks.
uv run train --updates 300 --name ppo
uv run evaluate runs/ppo/final.npz
train runs PPO on GPU through MJX, checkpointing to runs/<name>/, and resumes with
--resume runs/<name>/latest.npz. It writes TensorBoard events beside the checkpoints,
grouped into panels: task for return and tracking error, learning for the PPO
diagnostics, policy for action statistics, heldout for periodic held-out evaluation
against the naive tracker, strata for the per-mode and per-weight breakdown, and
speed for wall clock.
uv run tensorboard --logdir runs
``` `evaluate` scores a checkpoint against the offline solver at
eleven weights and prints the ratio of achieved cost to the best achievable cost.
uv run evaluate runs/ppo/final.npz --gui
`--gui` watches the policy instead of scoring it. The commanded target is drawn as a translucent
green arm, what the policy actually commands as a translucent yellow one, and two spheres mark the
arm centre of mass where the episode started and where it is now. Each target is replayed at
weights 0, 0.3, 0.6 and 0.9 so the trade-off is visible on the same reach; `--alpha` pins one, and
`--side left` or `--side right` targets a single arm, leaving the other free to counterweight.
## Status
The model, position-servo actuators, analytic kinematics, an offline solver, a trained policy, and a viewer entrypoint.
Each joint is driven by a `<position>` servo taking a target angle, so a joint holds its commanded pose against gravity and moves only when commanded. Control ranges follow the joint limits, and torque is capped at 120 N m per shoulder and 60 N m per elbow.
The joint limits are certified collision-free, so contacts are disabled. `geometry.py` provides forward kinematics, the arm centre of mass and its Jacobian, and an exact self-collision predicate, all matching MuJoCo to machine precision. `frontier.py` solves the trade-off between putting the bone ends where a target pose would put them and holding the arm centre of mass near where it started; a single weight selects an operating point.
One policy spans the whole trade-off, taking the weight as an input. A target may cover both arms
or only one: NaN entries in a target mark joints the task does not score, and the unscored arm is
free to move however it likes, typically counterweighting the scored one. Measured against the
offline solver on held-out targets, the policy reaches within about 10 percent of the best
achievable cost across the middle of the weight range.