GMTrack

Highly dynamic humanoid whole-body motion tracking on the Unitree G1 (29 DoF), packaged as an external task package for mjlab.

Two-stage continual learning (PACE + STAR) on top of the mjlab/BeyondMimic tracking stack, referencing Extreme-RGMT, SONIC and InstinctLab.

mjlab already ships a BeyondMimic reimplementation under mjlab/tasks/tracking/, so GMTrack plugs into mjlab instead of forking it: tasks register through mjlab's entry-point hook, and the learning code subclasses rsl-rl.

Method

D — all motion data (paper Table IV)
│
├─ Stage I ─ plain PPO + adaptive bin sampling over all of D ──────────────► π_base
│
├─ Stratification (runs only after Stage I)
│     sequences > 10 s are cut into 10 s logical segments, short ones kept whole;
│     π_base is rolled out 5× per segment under randomization
│     completion ≥ 80% → D_m (mastered)          otherwise → D_c (challenging)
│
└─ Stage II ─ ONE training run, two environment groups, one shared network
      acquisition   (ξ = 0.8): samples D_c, adaptive sampling, PPO + STAR,
                               domain randomization / pushes / observation noise
      consolidation (1 − ξ)  : samples D_m, uniform sampling, nominal dynamics,
                               distillation toward the frozen π_ref = π_base
                                                                          ► π_aug

• PACE — asymmetric acquisition/consolidation environment roles plus a progress-adaptive regularizer λ_con toward the frozen base policy, so learning new highly dynamic skills does not erase mastered ones.
• STAR — advantage-prioritized resampling of trajectory fragments drawn from the hardest temporal bins, which compensates for the sparse, failure-heavy samples that extreme motions produce.
• Policy — a three-branch encoder (proprioceptive state, past actions, command window) feeding a causal attention encoder with a finite-scalar-quantization bottleneck. The actor consumes only deployable observations.

Status

Research code under active development.

• The implementation is complete and covered by unit tests, and one end-to-end chain (Stage I → stratification → Stage II) has run to completion.
• No benchmark numbers are published yet — the five-seed sweep over the four test sets has not been run.
• Task configs, manifests, and interfaces may still change between commits.

Installation

Requirements:

• Linux x86_64 with an NVIDIA GPU — mjlab simulates through MuJoCo Warp, which has no CPU backend
• Python 3.10–3.13
• uv

git clone https://github.com/cmjang/GMTrack.git
cd GMTrack
uv sync
uv run list-envs | grep GMTrack

uv sync resolves torch from the CUDA 12.8 index. uv run re-syncs the environment on every invocation; call .venv/bin/python -m ... directly to skip that.

third_party/rsl_rl is a local editable checkout of rsl-rl-lib 5.4.0, the version mjlab pins exactly. It is never patched — every modification lives in src/gmtrack/rsl_rl/ and is injected through rsl-rl's "module:Class" resolution.

Pipeline

1. Prepare motions

prepare_motions writes one complete sequence per source CSV — no 10 s slicing happens here.

uv run python -m gmtrack.scripts.prepare_motions \
    --input-dir data/datasets/raw/lafan1 --source lafan1 \
    --input-format mjlab --input-fps 30 \
    --output-dir data/datasets/stage1_full/lafan1 \
    --manifest logs/data_build/manifests/lafan1_full.json

High-dynamic proxy clips are additionally aligned against the G1 collision geometry and must pass the clearance gate before entering a manifest:

uv run python -m gmtrack.scripts.prepare_motions \
    --input-dir data/datasets/raw/seed_stunts --source seed-stunts \
    --input-format bones-seed --input-fps 120 --output-fps 50 \
    --output-dir data/datasets/seed_stunts_grounded \
    --manifest logs/data_build/manifests/seed_stunts_grounded.json \
    --ground-alignment g1_collision --ground-clearance-m 0.003 \
    --ground-smoothing-radius-s 0.3 --device cuda:0

uv run python -m gmtrack.scripts.audit_ground_clearance \
    --manifest logs/data_build/manifests/seed_stunts_grounded.json --threshold 0.003

Replay the reference with a zero-action policy to sanity-check a manifest:

uv run play GMTrack-Stage1-Flat-Unitree-G1 --agent zero

2. Stage I

Environment count and iteration budget are baked into the task configuration, so no scale flags are needed:

uv run train GMTrack-Stage1-Flat-Unitree-G1

3. Stratification

Use the same manifest Stage I trained on, otherwise D_m ∪ D_c no longer covers the training distribution.

uv run python -m gmtrack.scripts.stratify \
    --checkpoint logs/rsl_rl/gmtrack_stage1/<run>/model_99999.pt \
    --manifest data/current/<stage1-manifest>.json

This produces stratified.json, mastered.json, challenging.json and stratification_report.json. The four files share one provenance hash and are validated as a set; hand-edited splits fail closed. Against paper Table V, D_c should land around 10% of the corpus — close to half means Stage I is undertrained.

4. Stage II

uv run train GMTrack-Stage2-Flat-Unitree-G1 \
    --agent.algorithm.base-checkpoint logs/rsl_rl/gmtrack_stage1/<run>/model_99999.pt

The same checkpoint warm-starts π_θ and freezes π_ref. A short gate run can be checked before committing to the long job:

uv run python -m gmtrack.scripts.check_stage2_gate --run-dir logs/rsl_rl/gmtrack_stage2/<run>

5. Evaluate and aggregate

uv run python -m gmtrack.scripts.evaluate \
    --checkpoint <ckpt> --manifest <test-set>.json \
    --training-seed 42 --out logs/eval/seed42.json

uv run python -m gmtrack.scripts.aggregate_evaluations \
    --inputs logs/eval/seed4{2,3,4,5,6}.json --out logs/eval/five_seed_summary.json

python -m gmtrack.scripts.webplay opens a Viser viewer through the evaluation harness for interactive inspection of a checkpoint; it loads policy weights only, never the training motion sampler.

Registered tasks

Task ID	Description
GMTrack-Stage1-Flat-Unitree-G1	Stage I, paper-default command window
GMTrack-Stage2-Flat-Unitree-G1	Stage II, paper-default command window
GMTrack-Stage1-Heading-Flat-Unitree-G1	Stage I with heading feedback
GMTrack-Stage2-Heading-Flat-Unitree-G1	Stage II with heading feedback
GMTrack-Stage1-Causal-Flat-Unitree-G1	Stage I, past-only command window
GMTrack-Stage1-Causal-Heading-Flat-Unitree-G1	Stage I, past-only window with heading feedback

Causal tasks use only current and past command tokens. Heading appends the 6D robot-to-reference pelvis orientation error to each command token.

Motion data

The paper trains Stage I on 3.096 h of retargeted 50 Hz motion:

Source	Paper duration	Availability
LAFAN1	2.444 h	public — lvhaidong/LAFAN1_Retargeting_Dataset, drop the CSVs into data/datasets/raw/lafan1/
AMASS	0.511 h	needs your own retargeting to G1
In-house Xsens captures	0.141 h	not public

Because the Xsens high-dynamic slice is unavailable, the registered Stage-I task defaults to a duration-matched proxy built from public high-dynamic sources (screened MotionDecode and BONES-SEED takes plus representative cartwheels). Any experiment report using it must say so — it is not the paper's data distribution.

Motion payloads are ignored by git. Point the registry at a local dataset without touching source code:

export GMTRACK_DATA=/abs/path/to/data              # dataset root
export GMTRACK_STAGE1_MANIFEST=/abs/path/to/stage1_complete_sequences.json
export GMTRACK_STRATIFIED_MANIFEST=/abs/path/to/stratified.json
export GMTRACK_MASTERED_MANIFEST=/abs/path/to/mastered.json
export GMTRACK_CHALLENGING_MANIFEST=/abs/path/to/challenging.json

Two conversion rules:

• Body ordering follows MuJoCo's depth-first convention. IsaacLab-derived converters silently track the wrong links, so use prepare_motions (or mjlab's csv_to_npz).
• Proxy clips retargeted from other skeletons are ground-aligned upward against all 29 active G1 collision geoms with a 3 mm clearance gate, re-differentiating root velocity afterwards. That threshold is a local data-cleaning choice, not a paper parameter.

Evaluation protocol

Following Sec. VI-A, the only failure criterion is root height deviating from the reference by more than 0.2 m. Each rollout runs from the reference's first frame to its end or to failure, and metrics are Succ. / E_MPJPE (mm, root-relative) / d_vel (mm/frame) / d_acc (mm/frame²), averaged over five training seeds.

The paper's four test sets are all distinct from D_c — D_c is a training set, and numbers on it cannot be compared to Table VI. Extreme-RGMT's published values, as Succ. / E_MPJPE:

Category	Test set	Stage I	Full (Stage II)
Generalist	In-source Motion (held out from LAFAN1 + AMASS)	99.54% / 40.07 mm	99.76% / 40.79 mm
Generalist	Unseen Motion (independent video reconstructions)	95.13% / 45.80 mm	96.68% / 46.91 mm
Specialist	XtremeMotion (public OmniXtreme high-dynamic set)	21.42% / 46.72 mm	100.00% / 40.18 mm
Specialist	AMASS Challenging (hard AMASS motions, direct retarget)	18.18% / 55.17 mm	90.91% / 46.39 mm

The specialist gap is what Stage II is for.

Checkpoints

Pretrained weights are published as GitHub Release assets rather than committed to git history:

Stage	Variant	Files
Stage I	no heading	stage1_no_heading_model_116500.pt
Stage I	heading	stage1_heading_model_99999.pt
Stage I	causal + heading	stage1_causal_heading_model_99999.pt
Stage II	no heading	model_17000.pt, model_31500.pt, model_32000.pt

Checkpoints produced before the current round of alignment fixes are fine for inspection but not for final numbers: retrain Stage I, rerun stratification, then retrain Stage II.

Development

uv run pytest tests/ -v
uv run ruff check src tests

Conventions:

• mjlab style: 2-space indent, ruff ["E4", "E7", "E9", "F", "I", "B"].
• Subclass rsl-rl rather than editing third_party/rsl_rl.
• Prefer raising over defensive fallbacks. A silent clamp or getattr(..., default) hides the class of bug that is hardest to find here: a plausible but wrong reference frame.
• Every paper constant carries a comment naming the table or equation it comes from.

Project layout

Path	Responsibility
src/gmtrack/pace.py	acquisition/consolidation environment split
src/gmtrack/mdp/	motion library, adaptive sampling, commands, observations, rewards, terminations, events
src/gmtrack/envs/env_cfg.py	full environment configuration (paper Tables I–II)
src/gmtrack/rsl_rl/models.py	three-branch encoder, causal attention, actor/critic
src/gmtrack/rsl_rl/fsq.py	finite scalar quantization bottleneck
src/gmtrack/rsl_rl/ppo_pace.py	PACE objectives and adaptive λ_con
src/gmtrack/rsl_rl/storage.py	STAR normalization and fragment resampling
src/gmtrack/motion_grounding.py	collision-geometry ground alignment
src/gmtrack/scripts/	data preparation, stratification, evaluation, viewers

References

• Extreme-RGMT: Continual Learning of Highly Dynamic Skills for Robust Generalist Humanoid Control — arXiv:2607.20110 — and its predecessor RGMT — arXiv:2601.23080.
• InstinctLab and SONIC.
• mjlab and rsl-rl, which this package builds on.

License

No license has been chosen yet, so default copyright applies and no reuse rights are granted. Open an issue if you need one.