WBC-MJLab — unified WBC motion tracking on mjlab (one policy, many motions)

[simeon-ned]

Hi @kevinzakka and mjlab folks,

Thank you for mjlab — it's become a big part of my day-to-day robotics work. The sim stack, task registration, and overall workflow make experimenting with humanoid control much easier and more pleasant than juggling yet another bespoke codebase.

I've been using it to pull together recent whole-body tracking ideas into one place, and wanted to share that effort: wbc-mjlab — an extension I've been building on top of mjlab.

A lot of recent humanoid WBC work (ZEST, BeyondMimic, SONIC, …) still ships as separate stacks. I wanted one training surface: a shared motion-tracking MDP where paper-specific choices are --task switches (RSI, observations, rewards, DR) — and on deploy, one policy for all motions, swapping clips at runtime instead of retraining per skill.

If this direction is interesting to you — comparing methods fairly, scaling to clip libraries, or pushing toward deploy — I'd genuinely love to hear from the community. Feedback, experiments, criticism, issues, and PRs are all very welcome.

What's in there already

Even though this is early, I've already pulled in a lot from recent WBC papers and wired it onto one shared MDP:

• ZEST-style tracking — ewards, reward-aligned RSI, assistive wrench, Zest terminations, per-joint action scales + DR
• BeyondMimic-style sampling — binary-failure adaptive RSI, whole-body keybody similarity, full-obs actor setup
• Shared infrastructure — EMA similarity metrics, motion command + playback, multi-clip datasets, deploy-style obs export
• Motion pipeline rework — GMR PKL ingest, batch GPU conversion to NPZ, optional motion-bundle cache
• Deploy export — policy.onnx + config.yaml from train/play
• Optional G1 runtime — wbc-g1-deploy: one ONNX policy, clip library via manifest.yaml

On that single deployed policy, the robot can already walk, jog, run, crawl, fight, get up from the floor, lie down, and do flips — clip swap at runtime, no checkpoint change.

Easy to configure your own WBC variant

The idea is tasks, not forks. Each paper's knobs live in a small env builder — see the G1 configs:

src/wbc_mjlab/robots/g1/configs

• zest.py — Zest paper repro
• binary_failure.py — BeyondMimic-style RSI
• wbc.py — default / deploy-oriented stack

Add a method → add a config module → register one WbcTaskConfig entry. Same CLI, same logs layout, comparable runs. Happy to help anyone who wants to plug in their robot or reproduce a specific paper setup.

Quick try (no big dataset download)

13 bundled sample CSVs (LAFAN1 + BONES-SEED excerpts):

git clone https://github.com/wbc-mjlab/wbc-mjlab.git && cd wbc-mjlab
uv run wbc-mjlab-data-to-npz --robot g1 --dataset samples
uv run wbc-mjlab-train --task Wbc-G1 --dataset samples
uv run wbc-mjlab-play --task Wbc-G1 --dataset samples

Same uv workflow as mjlab. Requires mjlab ≥ 1.4 + GPU for training.

Deploy demo — one policy, many clips

Same unified policy — sim loopback + real robot (clip swap at runtime):

Sim (loopback)
Walk wbc_g1_walk.mp4	Fight wbc_g1_fight.mp4	Jog wbc_g1_jog.mp4

Coming soon: side flips, backflips, and more high-dynamic clips on the same policy.

What's next (near term)

• Tech report
• Proper documentation (beyond the README landing page for now)
• More deploy demos and polish on the sim→real path
• SONIC-style task, CI, PyPI, …
• Pluggable robots & papers — a dedicated structure so you can add your robot or your paper's method as an external module/repo, without forking the core MDP

Status

Still early — APIs and tasks will move. Paper repros are starting points, not bit-for-bit guarantees.

If you're working on motion tracking in mjlab — whether you want to compare RSI schemes, bring a new robot, or sanity-check my design choices — please reach out. Issues, discussions, and PRs are exactly what I need right now.

[kevinzakka]

This is awesome, so great to hear @simeon-ned! Thanks for sharing :)

[simeon-ned]

We’ve added a small MuJoCo WASM web demo:) Feel free to try it directly in the browser!

BTW, @kevinzakka, would you mind if I open a PR adding a reference to our project under the project list in the mjlab docs?

[kevinzakka]

I saw, amazing job! Loving all the open source and democratization y'all are doing. Regarding the PR, go for it!

[simeon-ned]

Done, once you have time, could you please take a look?