Newton G1 Loco-Manipulation Training

Train a Unitree G1 humanoid (29 DoF) to walk to a 3D target and reach its right end-effector to the target, using the Newton physics engine via MuJoCo Playground and Brax PPO. Sim-only demo, single GPU.

Quick start

# One-time setup
./scripts/setup_env.sh
source .venv/bin/activate
./scripts/fetch_g1_asset.sh

# Stage 1 — locomotion baseline (~3–6 h on RTX 4090, slower on smaller GPUs)
python -m training.train --config configs/stage1_locomotion.yaml
python -m training.eval  --config configs/stage1_locomotion.yaml \
    --checkpoint runs/stage1_locomotion/final.pkl

# Stage 2 — standing reach (~2–3 h)
python -m training.train --config configs/stage2_reach.yaml \
    --init-checkpoint runs/stage1_locomotion/final.pkl
python -m training.eval  --config configs/stage2_reach.yaml \
    --checkpoint runs/stage2_reach/final.pkl

# Stage 3 — loco-manipulation (~9–12 h)
python -m training.train --config configs/stage3_loco_reach.yaml \
    --init-checkpoint runs/stage2_reach/final.pkl
python -m training.eval  --config configs/stage3_loco_reach.yaml \
    --checkpoint runs/stage3_loco_reach/final.pkl

A smoke config configs/stage1_smoke.yaml runs Stage 1 for only 10M steps (~30–60 min) to verify the stack.

Requirements

• Python 3.10+
• NVIDIA GPU (Maxwell or newer), driver 545+, CUDA 12
• Linux (x86-64 or aarch64). Newton's macOS build is CPU-only.

Repo layout

Path	Purpose
envs/g1_locomotion.py	Stage 1: wraps Playground's G1JoystickFlatTerrain with impl='warp'
envs/g1_reach.py	Stage 2: adds 3D reach target + extra rewards/obs
envs/g1_loco_reach.py	Stage 3: extends Stage 2 with far target + internal command generator
envs/common/rewards.py	Pure reward primitives (reach_distance, reach_bonus, arm_action_rate)
envs/common/observations.py	Pelvis-frame transforms
training/train.py	Brax PPO entry, reads YAML config, supports --init-checkpoint
training/eval.py	Deterministic rollouts + metrics + MP4 video
training/checkpoint.py	Pickled-param save/load +partial_restore for cross-stage init
configs/	Per-stage YAML configs (full + smoke variants)
assets/g1/	Unitree G1 MJCF (fetched by scripts/fetch_g1_asset.sh, not committed)
runs/	Training outputs (logs, checkpoints, videos) — gitignored

After fetching the G1 asset

Verify body/site names in envs/g1_reach.py match the MJCF:

source .venv/bin/activate
python -c "
import mujoco
m = mujoco.MjModel.from_xml_path('assets/g1/scene.xml')
print('bodies:', [m.body(i).name for i in range(m.nbody)])
print('nu (actuators):', m.nu)
"

If PELVIS_BODY or RIGHT_EE_BODY in envs/g1_reach.py don't match, edit those constants. If nu != 29, also update RIGHT_ARM_INDICES.