mjlab-x

Setup

Clone mjlab, thats it!

git clone git@github.com:wuphilipp/mjlab-x.git

Velocity Training

MUJOCO_GL=egl CUDA_VISIBLE_DEVICES=0 uv run train \
    Mjlab-Velocity-Flat-Xdof-X0 \
    --video True --video-interval 10000 \
    --env.scene.num-envs 4096 \
    --agent.max-iterations 50_000

For multi GPU add --gpu-ids all, and divide num-envs by the total number of GPUs (since each GPU runs the full num-envs count).

X20 Height Control

The height command is an absolute squat amount: 0 is stand and 1 is squat.

MUJOCO_GL=egl CUDA_VISIBLE_DEVICES=0 uv run train \
    Mjlab-Height-Absolute-Xdof-X20 \
    --env.scene.num-envs 4096 \
    --agent.max-iterations 30000

Symmetric arm + squat control

Mjlab-Height-Arm-Xdof-X20 exposes one five-dimensional command:

[squat, shoulder_pitch, shoulder_roll, shoulder_yaw, elbow]

All height commands use a squat amount in [0, 1]: 0 is standing and 1 is the deepest squat. This convention is shared by the scalar height, height+velocity, and height+arm tasks. The four arm values are in [-1, 1] and are mirrored onto both arms. All five zeros therefore request standing with the first four joints of both arms at physical joint angle zero. Positive shoulder_roll raises both arms away from the torso; shoulder_roll=1 is the straight-arm T-pose anchor. Random training commands cover arbitrary combinations, with extra sampling of neutral, T-pose, and bent-elbow curl anchors. New targets are sampled every 2–4 seconds within each episode, and the five command values themselves move toward those targets at bounded rates so the policy learns smooth transitions rather than responding only to command steps. Training also randomizes an independent 0–3.5 kg payload on each terminal hand link. Push strength ramps smoothly from the standing-task baseline to 1.5× between iterations 12k and 18k.

MUJOCO_GL=egl CUDA_VISIBLE_DEVICES=0 uv run train \
    Mjlab-Height-Arm-Xdof-X20 \
    --env.scene.num-envs 4096 \
    --agent.max-iterations 30000 \
    --video True --video-interval 5000

The height-only task resets across uniformly sampled heights. Squat amount 0 is the 0.81 m stand, approximately 0.7222 is the tuned 0.55 m intermediate squat, and 1 is the tuned 0.45 m deep squat. Pose and root-orientation targets interpolate through those height-aligned waypoints.

To train velocity and height control in one policy, use the combined task. Its command observation is [vx, vy, yaw_rate, height_level]. For this locomotion policy, squat amount 0 is the 0.81 m stand and 1 is the safer 0.55 m intermediate squat; the 0.45 m deep squat is reserved for the height-only task.

MUJOCO_GL=egl CUDA_VISIBLE_DEVICES=0 uv run train \
    Mjlab-Height-Velocity-Xdof-X20 \
    --env.scene.num-envs 4096 \
    --agent.max-iterations 30000 \
    --video True --video-interval 5000

Tune the X20 squat pose

Launch the browser-based pose editor, then open http://localhost:8099:

uv run tune-x20-squat

The sliders mirror the left/right leg joints, update the robot live, and show a green STAND_KEYFRAME reference. The generated Python keyframe is displayed in the UI; the Save button writes the selected values to /tmp/x20_squat_pose.json by default.

Motion Imitation (BeyondMimic-style Tracking)

Train the X15 to imitate retargeted motion-capture clips. The entire workflow runs off a single piece of state — the wandb run path: you train here, hand someone the run path, and they deploy in humanoid_internal with nothing else.

S3 v15_motion library ──prepare-motion──▶ wandb motions registry ──train──▶ run
   (xmjlab only)                                                              │
                                                  deploy --run xdof/mjlab/<id>┘
                                              (humanoid_internal, wandb-only)

The motion library lives at s3://xdof-internal-research/humanoid/v15_motion/ (source_lafan/, retarget_omni/, tracking/, preview/, manifest.json). xmjlab pulls from S3; humanoid_internal never touches S3.

1. Pick a motion and push it to the registry

prepare-motion pulls the tracking clip(s) from S3, applies any frame trimming, and uploads motion.npz + config.json to the wandb motions registry. It prints the --registry-name to train with.

# Single clip
uv run prepare-motion --clips walk1_subject1

# Frame-trimmed (start:end), or a named multi-clip set
uv run prepare-motion --clips walk1_subject1:100:600
uv run prepare-motion --clips walk1_subject1 run1_subject2 --name x15-walk-run

2. (optional) Visualize before training

# Pull a clip from S3 and view it
uv run visualize-clip --clip walk1_subject1

# Render headless to video
MUJOCO_GL=egl uv run visualize-clip --clip walk1_subject1 --video /tmp/preview.mp4

3. Train

MUJOCO_GL=egl uv run train \
    Mjlab-Tracking-Flat-Xdof-X15 \
    --registry-name xdof/csv_to_npz/walk1_subject1 \
    --env.scene.num-envs 4096 \
    --gpu-ids all \
    --agent.max-iterations 30000 \
    --video True --video-interval 5000

mjlab downloads the motion from the registry and links the artifact to the run (via use_artifact), so the run path alone is enough to recover the exact motion later — that is what makes deploy self-contained.

4. Deploy (in humanoid_internal, wandb-only)

# On the deploy-mjlab-x15 branch — only the run path is needed.
# macOS needs mjpython for the viewer; Linux with a display can use python.
mjpython -m xhum.sim.deploy_mujoco --config-path xdof/mjlab/{run-id}

Deploy downloads the ONNX policy and resolves the reference motion from the run's linked motions artifact. No S3, no pre-staged files.

Local / debug path

For iterating on a clip that isn't in the library yet, the local path still works end to end:

# Convert a holosoma qpos NPZ -> 50 fps tracking NPZ
uv run convert-motion --input ~/lafan_holosoma/holosoma_npz/walk1_subject1.npz \
    --output /tmp/walk1.npz                       # or --input-dir/--output-dir for batch

# Build motion.npz without uploading
uv run prepare-motion --clips walk1_subject1 --no-upload --output /tmp/motion.npz

# Train / play straight from a local file
uv run train Mjlab-Tracking-Flat-Xdof-X15 --env.commands.motion.motion-file /tmp/motion.npz ...
uv run play  Mjlab-Tracking-Flat-Xdof-X15 --wandb-run-path xdof/mjlab/{run-id} --motion-file /tmp/motion.npz

Available clips

Run uv run python -c "from xmjlab.motion import MotionStore; print(MotionStore().list_clips())" for the live list, or read manifest.json. The library is the full LAFAN set retargeted to the X15 (walk / run / jump / sprint / dance / fight / aiming / …).

Notes

• The task uses mjlab's BeyondMimic implementation (make_tracking_env_cfg) with adaptive reference state initialization.
• Actor observations exclude motion_anchor_pos_b and base_lin_vel (no state estimation on the real robot yet). The critic still sees them.
• The "humanoid" body must be first in body_names because MotionCommand._resample_command uses body_names[0] as the free-joint root. On the X15 the free joint belongs to "humanoid", not "pelvis".
• Multi-clip sets are concatenated; each join is a 1-frame velocity discontinuity (logged). Single-clip is the canonical one-policy-per-motion path.

Evaluation

To evaluate your trained policy, you'll need your wandb run path. You can find this in the run overview. It follows the format {your-organization}/{project-name}/{run-id}, where run-id is a unique 8-character identifier.

uv run play \
    Mjlab-Velocity-Flat-Xdof-X0 \
    --wandb-run-path {wandb-run-path} \
    --num-envs 8