WBC-Mjlab: Whole Body Control in MuJoCo Lab

One shared MDP for whole-body motion tracking on mjlab — train once on a motion library, deploy one policy for many skills.

Recent works (ZEST, BeyondMimic, SONIC, OmniXtreme) is all WBC / large-scale tracking, with overlapping ideas (keybody rewards, adaptive sampling, multi-clip training) but different design choices — each still tends to ship as its own codebase. In wbc-mjlab, paper-specific knobs are --task switches on a shared stack:

• Multi-motion by design — train on multi-clip datasets (LAFAN, SEED, custom NPZ libraries); one controller generalizes across the library. At runtime, provide reference - via playing clip, teleop, or higher level policy.
• Shared MDP — rewards, terminations, motion command, RSI, and playback live in env/ once; robots and papers plug in via task configs.
• Tasks, not forks — ZEST, BeyondMimic-style RSI, deploy obs, etc. are --task switches (Wbc-G1, Wbc-G1-Zest, Wbc-G1-BinaryFailure, …) with the same CLI and log layout for fair comparison.
• Motion data pipeline — versioned libraries under data/, GMR PKL ingest, batch GPU CSV→NPZ, optional motion-bundle cache.
• Building blocks — paper presets in presets/ and G1 task builders in robots/g1/tasks.py; add a paper setup or tune your own WBC without forking the core MDP.
• Plug-in robots — external packages register via register_wbc_extension (same MDP, same CLIs); see wbc-mjlab-extension-h2 for a reference layout.
• One policy, many skills — one policy for walk, jog, run, crawl, fight, get-up, lie-down, flips, and more.
• Sim → real — train/play export policy.onnx + config.yaml aligned with the deploy runtime.

Details: Documentation · CONTRIBUTING.md

Quick start

Requires mjlab (≥ 1.5) and an NVIDIA GPU for training.

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

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

Convert trajectory samples (13 source CSVs manifest & credits) to npz - calculating FK for body targets, velocities etc:

uv run wbc-mjlab-data-to-npz --robot g1 --dataset samples --batch-size 8

Demo — live web demo (trained policy in the browser); local play with bundled checkpoint (convert samples first):

uv run wbc-mjlab-demo

Train on converted npz library (check mjlab train args for resuming, number of envs etc):

uv run wbc-mjlab-train --task Wbc-G1 --dataset samples

Evaluation of last exported log on library (check args for viewer, choosing chekpoint, motion etc):

uv run wbc-mjlab-play --task Wbc-G1 --dataset samples

Related repos