wbc-mjlab

Unified Whole-Body Motion Tracking on mjlab

Simeon Nedelchev¹ Anton Chaplygin² Lev Kozlov³ Ivan Domrachev³

¹Innopolis University ²Robotics Center ³KAIST

[Code] [Docs] [Demo] [PyPI] [Video] [Colab]

Unitree G1 — one policy, many skills (get-up, idle, dance, fight, sprint, sideflip).

Abstract

We present wbc-mjlab, an open-source training library for whole-body motion tracking on mjlab. The library factors the low-level tracker shared by recent humanoid systems into one manager-based MDP: scaled residual joint actions, a multi-clip motion command with adaptive reference-state initialization (RSI), exponential tracking rewards on anchor and keybody kinematics, and a composable catalog of motion-reference observations. Paper-specific choices—ZEST-style end-effector tracking, BeyondMimic-style binary-failure RSI, state-estimation actor layouts—are expressed as presets on the same template rather than separate codebases. Robots plug in via extension packages; training emits artifacts for modular deploy (hardware runtime and in-browser demo).

Interactive demo

A real MuJoCo sim and ONNX whole-body policy run live in your browser (WASM — no install). The robot tracks a reference motion clip; (open fullscreen for the full UI, including generator mode).

Controls

pick a clip from the dropdown (top right).
play, pause, loop clip (not policy).
get up when down, lie down when idle.
click-drag the robot to push/pull; empty space orbits the camera.
reset pose.

Try it fast

Fastest path: try the live demo (no GPU). To train locally you need mjlab (≥ 1.5) and an NVIDIA GPU. Full install notes are in the docs.

Browser demo. No install — open the WASM + ONNX gallery and switch clips. wbc-mjlab.github.io/wbc-demo

Clone & list tasks.

git clone https://github.com/wbc-mjlab/wbc-mjlab.git && cd wbc-mjlab
uv run wbc-mjlab-list-envs

Convert sample clips → NPZ. Bundled G1 samples (CSV → FK bodies / velocities).
```
uv run wbc-mjlab-data-to-npz --robot g1 --dataset samples --batch-size 8
```
Train (default WBC). Paper recipes are --task switches on the same MDP (e.g. Wbc-G1-Zest).
```
uv run wbc-mjlab-train --task Wbc-G1 --dataset samples
```
Play / export. Play disables training aids so rollouts match deploy; exports policy.onnx + config.yaml.
```
uv run wbc-mjlab-play --task Wbc-G1 --dataset samples
```

uv run syncs from uv.lock on first use. CUDA/CPU PyTorch: make sync / make sync-cpu. See also the installation guide.

Particularities

Config assembly. Each task is base MDP template \(\rightarrow\) robot entity \(\rightarrow\) preset(s) \(\rightarrow\) registered --task id. Presets are pure config functions on one template (no env fork); ablations are task switches, not code branches. Adaptive RSI samples a (clip, temporal bin, frame) on episode reset inside the motion command—not between reward and observation each step.

Single command, slot-based MDP. Rewards, RSI, reference observations, and assistive wrench all read one motion command. Parallel envs may track different clips from the same NPZ library; body lists and weights stay in config, not hard-coded skeletons.

Scaled residual actions. The default policy outputs a residual \(a_t\) added to the motion-command reference joints before PD. Per-joint scales \(s\) travel with the robot entity (often \(s_j \approx 0.25\,\tau^{\max}_j / k_{p,j}\)), so the same preset stack ports across humanoids.

\[ q^{\mathrm{des}}_t = q^{\mathrm{ref}}_t + s \odot a_t \]

Anchor-relative keybody targets. Style rewards compare limbs in a yaw-aligned anchor frame (horizontal from the robot anchor, height from the reference). Matches deploy where global drift is partly unobservable but relative limb pose is not.

Adaptive RSI. Clips are partitioned into bins of width \(\Delta \approx 4\,\mathrm{s}\); failure EMAs drive softmax sampling with a uniform floor \(\rho \approx 0.15\). Default WBC/ZEST score bins with the same weighted tracking rewards used for learning (reward-aligned); a BeyondMimic-style preset switches to binary-failure RSI (early terminate \(\Rightarrow\) \(f = 1\)).

\[ \bar f_{\tau,b} \leftarrow (1-\alpha)\,\bar f_{\tau,b} + \alpha\, f_{\tau,b} \]

Modular reference observations. Reference kinematics are independent catalog terms (not one bundled command vector). Presets drop, swap, or reorder entries—and add per-term noise—without changing playback or rewards. A state-estimation actor preset composes with any tracking preset by swapping pose-error terms into the actor.

Assistive wrench (train only). Optional ZEST-style spatial wrench on the anchor with gain \(\beta \in [0, \beta_{\max}]\) from the sampled bin’s failure level. Play and deploy disable it—it is a curriculum aid, not part of the exported policy.

Train vs. play split. Play disables observation corruption, motion DR, assistive wrench, and pushes so in-sim evaluation matches deploy; training keeps those aids on.

In-tree G1 tasks. Paper knobs are configuration: Wbc-G1 (all keybodies, reward-aligned RSI), Wbc-G1-Zest (end-effector body rewards), Wbc-G1-BinaryFailure (BeyondMimic-style binary RSI), and …-SE variants that swap in a state-estimation actor (pose errors + base linear velocity).

Export contract. Training emits policy.onnx + config.yaml for modular runtimes (wbc-g1-deploy, wbc-demo) without importing the trainer. At deploy, the reference stream may be clips, teleop, or a higher-level policy; planning UIs and drivers stay outside the training repo.

Robot extensions. New platforms ship as packages via register_wbc_extension (entity, task builder, data paths)—same CLIs as G1. See wbc-mjlab-extension-h2; no changes to the shared tracker required.

Code

wbc-mjlab — Training — shared MDP, presets, G1 tasks, ONNX export
wbc-mjlab-extension-h2 — Reference H2 robot extension (no core fork)
wbc-g1-deploy — G1 onboard runtime — one policy, clip library
wbc-demo — Live MuJoCo + ONNX in the browser

BibTeX

@misc{wbc-mjlab,
  title        = {wbc-mjlab: Unified Whole-Body Motion Tracking on mjlab},
  author       = {Nedelchev, Simeon and Chaplygin, Anton and Kozlov, Lev and Domrachev, Ivan},
  year         = {2026},
  howpublished = {\url{https://github.com/wbc-mjlab/wbc-mjlab}}
}