booster_mjlab

booster_mjlab is an mjlab integration for the Booster K1. It provides the K1 robot model, velocity and motion-tracking tasks, and an Adversarial Motion Priors (AMP) training pipeline that learns natural-looking gaits from motion capture data.

Real-robot and simulation demos are on the project page, which also runs the velocity policy live in the browser (MuJoCo WebAssembly + the exported network).

Getting Started

booster_mjlab requires an NVIDIA GPU for training. macOS is supported for evaluation only.

Install from source:

# Install uv if you haven't already
curl -LsSf https://astral.sh/uv/install.sh | sh

git clone https://github.com/IntelligentRoboticsLab/booster-mjlab.git && cd booster-mjlab
uv run list_envs

The last command syncs the environment and prints every registered task.

Training Examples

1. Velocity Tracking

Train a Booster K1 to follow velocity commands on flat terrain. The reward has been tuned to work together with AMP reference motions. The default motion dataset contains a few locomotion clips retargeted from LAFAN1.

uv run train Mjlab-Velocity-Flat-Amp-DA-Muon-Booster-K1 --env.scene.num-envs 4096

Evaluate a policy while training (fetches latest checkpoint from Weights & Biases):

uv run play Mjlab-Velocity-Flat-Amp-DA-Muon-Booster-K1 --wandb-run-path your-org/mjlab/run-id

To train on your own motions, pass a Hugging Face dataset (namespace/repo) or a local motion file or directory with --agent.dataset-root.

The -DA- tasks train with left/right symmetry data augmentation. Drop DA from the task ID to train without it.

The -Muon- tasks use the Muon optimizer (Jordan et al., 2024) for the actor and critic weight matrices. Drop Muon from the task ID for the Adam variant with otherwise identical settings.

2. Motion Imitation

Train a Booster K1 to track motion capture clips on flat terrain. Motions are managed as WandB artifacts; see Preparing Motions for Tracking for converting and uploading GMR (.pkl) or BeyondMimic (.csv) clips.

Once a motion is uploaded, train and evaluate with:

uv run train Mjlab-Tracking-Flat-Booster-K1 --registry-name your-org/motion_upload/motion-tracking:v0 --env.scene.num-envs 4096
uv run play Mjlab-Tracking-Flat-Booster-K1 --wandb-run-path your-org/mjlab/run-id

For a motion prepared with csv_to_npz, pass --registry-name your-org/motions/motion_name instead.

3. Browse Motion Data

Inspect, play back and edit motion clips in the browser before training on them:

uv run visualize-motions --help

Clips from the lafan_locomotion_k1 dataset, retargeted onto the parallel-ankle K1:

Tasks

Velocity tasks come in Flat and Rough terrain variants, optionally with -Amp- (motion prior) and -DA- (symmetry data augmentation). Tasks suffixed -Parallel use the parallel-linkage ankle model of the K1. Tracking tasks run on flat terrain. List them all with:

uv run list_envs

Deployment

Example code for deploying trained policies on the real Booster K1 is coming soon.

Hungry for more?

booster_mjlab is a whIRLwind project. Our mission is to push the boundaries of robotics and AI in robot football. We are always looking for interested students or collaboration partners to join us. Learn more at whirlwind.team.