InstantAMP: Train AMP in 30 minutes on 1 RTX 5090

InstantAMP trains a Unitree G1 velocity-tracking policy with an adversarial motion prior (AMP) in about 30 minutes on a single RTX 5090. The policy walks, runs, side-steps and turns like the human clips it learns from, and it gets back up after a fall.

The policy is not trained with model-free RL alone. Like SuperMimic, InstantAMP fits a differentiable world model of the robot from simulator transitions. The policy is then optimized with analytic gradients through short imagined rollouts (backpropagation through time). The style rewards of the two motion priors are differentiable in the imagined states too, so the policy follows their gradient through the world model like any other reward term.

The whole training loop is compiled with JAX, including the MuJoCo Warp simulator, and runs on one GPU. The code builds on XSoma.

Contents

• What is trained
• Results
• Setup
• Motion data
• Training
• Flags
• Evaluation and rendering
• Repository layout
• Acknowledgements and citation
• License

What is trained

Part	What it is
Robot	Unitree G1, 29 joints. mjlab's MJCF with full collision (self-collisions included) and the Unitree actuator model (torque-speed curve and friction), 200 Hz physics, 50 Hz policy.
Task	Follow velocity commands on flat ground: vx ∈ [−1.5, 3.0] m/s, vy ∈ [−1, 1] m/s, wz ∈ [−π/2, π/2] rad/s. 25 % of the envs use heading control and 5 % stand still. Commands are redrawn every 3 to 8 s.
Observation	4 frames of history (384 values): pelvis angular velocity, gravity, command, joint positions and velocities, last action. Noise is added, and the policy normalizes the observation with running statistics.
Rewards	Velocity tracking on the pelvis, regularizers (joint acceleration, joint limits, action rate, foot slip, self-collision, wrist pose), a termination penalty, and a rise penalty while the robot is down.
Falls and get-up	An upright robot is terminated as soon as its pelvis, hips, shoulders, torso or hands touch the floor. 20 % of the resets start on the floor in a frame of a LAFAN1 get-up clip, and the robot has 5 s to stand up.
Robustness	Pushes every 5 to 10 s, shoves at reset, randomized foot friction, encoder bias and torso CoM.
Motion priors	Two WGAN-GP discriminators over 10-frame windows: one for walking and running (used while the robot is up) and one for getting up (used while it is down).
Learner	World model (MLP 3×512, 8-step open-loop loss), BPTT over a 48-step horizon (about 1 s), a critic that bootstraps the return at the horizon, an ODE actor, and a hybrid PPO term on 10 % of the real collector steps.

The reward that the policy maximizes in imagination is

r = (1 − amp_lerp) · r_task + amp_lerp · dt · (ramp · gate_loco · style_loco + gate_getup · style_getup)

• gate_loco = moving · (0.25 + track_lin) while the robot is up, so tracking well earns more style reward. It is 0 for a standing command.
• gate_getup = 1 while the robot has not been up yet in its episode.
• Each style is (1 + tanh((d − mean robot score) / 5)) / 2. Centering it on the robot's own mean score keeps the reward slope constant, however far the robot still is from the demos.

Results

The GIFs and curves above come from a 60-minute run on one RTX 5090 with 8192 environments. The GIFs show its final checkpoint. The curves are measured on 512 robots × 20 s per checkpoint, while the robots are up and commanded to move. Episodes are close to full length (1000 steps) from about 20 minutes on. The velocity error falls to about 0.45 m/s at 30 minutes and 0.35 m/s at 60 minutes.

The table below is the reference run: one RTX 5090 with 4096 environments and the default flags. Compilation took 146 s before training started, and training ran at 3.4 iterations per second.

Training time	Episode length (of 1000 steps)	Episodes ending in a fall	Linear tracking (per s)	Get-ups (rose / down starts)
5 min	85	1.00	0.008	80 / 485
10 min	152	1.00	0.021	100 / 117
15 min	810	0.23	0.40	34 / 35
20 min	949	0.07	0.70	40 / 40
30 min	935	0.10	0.73	27 / 33
60 min	993	0.01	0.80	20 / 20

The robot learns to get up first (by about 10 minutes) and walks from about 15 minutes on. We rolled out the checkpoint from about 55 minutes for 20 s in 256 environments. It ran upright and symmetrically, with its arms swinging at its sides, at commands up to 2.7 m/s. Its median velocity error was 0.17 m/s, and it fell once.

These numbers come from one seed. Expect some spread between seeds and GPUs.

Setup

You need Linux with an NVIDIA GPU, Python 3.11, and an NVIDIA driver with CUDA 12.8 support or newer. Warp's PyPI wheel needs driver 580 or newer; the reference run used 580.

git clone https://github.com/RoboticsResearchHubVN/InstantAMP.git
cd InstantAMP
python3.11 -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt

With an older driver, install Warp's CUDA 12 build (warp_lang-1.19.0+cu12) from its releases instead of the PyPI wheel.

Check the install on the CPU (no GPU or data needed):

JAX_PLATFORMS=cpu python -m pytest tests

Motion data

The clips are not part of this repository; their own licenses apply. One script downloads them and converts them to 50 Hz .npz files in the G1's joint order. It needs git, git-lfs and curl.

bash tools/download_data.sh

Folder	Clips	Used for	Source
data/amp/walk_run	30 clips, 156 s: walking, turning, side-stepping, running, running backwards, transitions	demos of the locomotion prior, upright reset frames	legged_lab (GMR retargets of AMASS / ACCAD)
data/recovery/lafan1	6 fallAndGetUp takes, 935 s (24 056 ten-frame windows in which the robot is down)	demos of the get-up prior, floor reset frames	Unitree's LAFAN1 retarget (LAFAN1 by Ubisoft La Forge, CC BY-NC-ND 4.0)

To use your own clips, convert them with tools/convert_motions.py and pass their folders with --amp_dir / --walk_dir (locomotion) or --amp_getup_dir / --floor_dir (get-up). The script reads legged_lab pickles (amp) or Unitree-style CSV files (csv: root position, root quaternion x y z w, then the 29 joints, 30 fps by default). An .npz clip holds fps (50), root_pos, root_quat (x, y, z, w), root_lin_vel, root_ang_vel (world frame), dof_pos and dof_vel (29 joints in the order of instantamp/env.py's DOF).

Training

python -m instantamp.train --log_dir runs/instantamp            # 30 minutes, the reference recipe

Every option of Config in instantamp/train.py is a command-line flag of the same name (--name value). Some common variations:

python -m instantamp.train --minutes 60 --log_dir runs/long                    # train longer
python -m instantamp.train --num_envs 2048 --log_dir runs/small                # less GPU memory
python -m instantamp.train --amp_lerp 0.3 --amp_style_scale 10 --log_dir runs/weak_style   # weaker style
python -m instantamp.train --amp_getup 0 --reset_floor 0 --log_dir runs/no_getup          # locomotion only
python -m instantamp.train --amp_getup_full 1 --imag_down 0.25 --log_dir runs/getup      # stronger get-up learning

A run writes the following to --log_dir:

• config.json: every option of the run.
• log.jsonl: one line every 50 iterations.
• ckpt/NNNN.npz and ckpt/NNNN.json: the averaged and live policies with the observation normalizer, every 2 minutes and at the end.
• ckpt/full/NNNN.npz: the whole training state.

The console line reads:

it 4300 | 20.13 min | 17.7M | len 949 fell 0.07 lin 0.700 ang 0.779 r 0.0362 | up 0.99 get-ups 40/40 | wm ... | G ... | 3.4 it/s
     | h_err 0.3/1.2 cm term env ... wm ... | amp d demo/robot r use ramp | getup d ... | V end ... | ppo clip ...

Field	Meaning
len	mean length of the finished episodes (1000 = a full 20 s)
fell	share of the finished episodes that ended by termination
lin, ang	tracking terms summed over an episode, divided by 20 s (1.0 is perfect tracking)
get-ups x/y	of the episodes that started on the floor, how many rose
wm, e8, h_err	world-model loss and its open-loop pelvis-height error after 8 / 32 steps
G, r_im, surv	imagined return, imagined reward per step, share of imagined rollouts that survive the horizon
amp d demo/robot	mean discriminator scores of the demos and of the robot; r is the AMP reward per step

The JAX compilation cache (--jax_cache, ~/.cache/instantamp/jax) makes later starts much faster.

Flags

All flags with their defaults (the reference recipe). Booleans are 0 / 1.

Environment

Flag	Default	Meaning
--num_envs	4096	parallel simulated robots
--nconmax, --njmax	64, 640	contacts and constraint rows per world (MuJoCo Warp pools)
--walk_dir	data/amp/walk_run	clips whose frames are the upright reset states
--floor_dir	data/recovery/lafan1	get-up clips whose down frames are the floor reset states
--reset_floor	0.2	share of resets that start on the floor
--prob_shove	0.2	share of upright resets that get a random kick of root and joint velocities
--push	1	overwrite the root's horizontal velocity with U[−1, 1] m/s every 5 to 10 s
--random_start	1	random initial episode counters, so that episodes do not all end together
--term_tau	0.02	width (m) of the sigmoid that gives imagined terminations their gradient

Motion priors (AMP)

Flag	Default	Meaning
--amp_w	2.0	> 0 turns AMP on; with --amp_lerp 0 it is the style weight
--amp_lerp	0.5	share of the style in the reward: (1 − lerp)·task + lerp·dt·style; 0: task + amp_w·dt·style
--amp_style_scale	5.0	scale of the style's tanh around the robot's mean score (smaller = stronger style); 0: midpoint / half-gap of the two mean scores
--amp_gate	track_shift	gate of the locomotion style: track_shift (moving × (lo + track_lin)), moving, track (moving × track_lin) or none
--amp_track_lo	0.25	the floor lo of the track_shift gate
--amp_dir	data/amp/walk_run	demos of the locomotion prior
--amp_getup	1	second discriminator for the get-up style
--amp_getup_dir	data/recovery/lafan1	demos of the get-up prior (windows in which the robot is down)
--amp_getup_full	0	1: the get-up style is at full scale from the first iteration, instead of sharing the ramp
--amp_frames	10	frames per discriminator window (0.2 s)
--amp_wrists	1	the wrist joints are part of the discriminator's frame
--amp_disc_steps	2	discriminator steps per iteration (each prior)
--amp_batch	4096	windows per discriminator step
--amp_recent	100	robot windows come from the newest this many buffer steps
--amp_lr	1e-4	discriminator learning rate
--amp_delay, --amp_ramp	500, 4000	the style reward starts after delay iterations and ramps up over ramp iterations

World model

Flag	Default	Meaning
--wm_hidden	512,512,512	hidden widths
--wm_window	8	open-loop steps of the loss
--wm_batches, --wm_batch	3, 2048	gradient steps per iteration and windows per step
--wm_lr	4e-4	learning rate
--wm_muon	1	Muon on the hidden weight matrices (Adam elsewhere)
--wm_gate	2.0	skip the policy update while the loss is above this times its running mean (0: off)
--wm_eval	1	log open-loop errors and termination rates

Policy

Flag	Default	Meaning
--horizon	48	imagined steps per rollout (about 1 s)
--pi_batch	2048	imagined rollouts per update
--pi_updates	2	policy updates per iteration
--pi_lr	3e-4	learning rate
--pi_clip	1.0	gradient-norm clip
--pi_gn_gate	10.0	skip an update whose gradient norm is above this times its running mean (0: off)
--pi_ema	0.995	decay of the averaged policy that checkpoints hold
--gamma	0.99	discount
--ode	1	ODE actor (ODE Policy): the network outputs the change of the action, integrated from the previous action
--ode_dmax, --ode_amax	3.24, 20.0	bounds of the action change per step and of the action (0: none)
--action_noise	0.3	exploration noise (std) of the noisy collector envs
--det_share	0.5	share of collector envs without noise (their windows train the discriminators)
--cmd_fresh	0.5	share of imagined rollouts that start with a freshly drawn command
--imag_down	0.0	share of imagined rollouts that start in a down state (0: uniform)
--obs_clip	10.0	clip of the normalized observation

Critic

Flag	Default	Meaning
--critic	1	value function that bootstraps the imagined return at the horizon (0: pure BPTT)
--critic_hidden	512,512,256	hidden widths
--critic_lr, --critic_epochs, --critic_mb	3e-4, 5, 4	TD(λ) fit on each batch of imagined rollouts
--lam	0.95	λ of the TD(λ) targets and of PPO's GAE

Hybrid PPO

Flag	Default	Meaning
--ppo_frac	0.1	share of the policy gradient's samples taken from real collector steps (0: off)
--ppo_len	8	newest steps of the PPO envs per batch
--ppo_clip	0.2	PPO ratio clip
--ppo_w	1.0	the PPO gradient is scaled to ppo_w times the norm of the BPTT gradient (< 0: p / (1 − p))

Schedule and runtime

Flag	Default	Meaning
--minutes	30.0	training time, without compilation
--env_steps	1	simulator steps of all envs per iteration
--buffer_steps	200	replay buffer length, in steps of all envs
--warmup_steps	32	steps collected before learning
--wm_only_iters	200	first iterations fit only the world model
--imag_unroll	4	unroll of the imagination scan
--imag_remat	none	dots or full: rematerialize imagined steps (less memory, slower)
--mm_bf16	imag	bf16 matmul outputs: imag, all or none
--log_every	50	iterations per log line
--ckpt_minutes	2.0	minutes between checkpoints
--seed	0	random seed
--log_dir	runs/instantamp	output folder
--jax_cache	~/.cache/instantamp/jax	JAX compilation cache (empty: off)
--xla_flags	--xla_gpu_enable_triton_gemm=false	extra XLA flags, comma-separated

JAX takes 30 % of the GPU's memory by default, and MuJoCo Warp allocates the rest at run time. Set XLA_PYTHON_CLIENT_MEM_FRACTION to change JAX's share.

Evaluation and rendering

python -m instantamp.evaluate runs/instantamp --last 3       # 1024 robots, 20 s, per checkpoint

For each checkpoint the evaluation prints and writes to eval.jsonl:

• the reward terms per second;
• the mean velocity error while up;
• the share of upright starts that fell;
• the share of floor starts that rose.

To render a walking GIF and a get-up GIF of a checkpoint (needs OpenGL; the 19 MB of visual meshes are downloaded once):

bash tools/download_meshes.sh
python tools/render.py runs/instantamp --out figures/instantamp     # figures/instantamp_walk.gif, _getup.gif

To check that imagination's reward terms match the simulator's on real transitions (needs a GPU):

python tools/check_heads.py

Repository layout

instantamp/
  env.py         robot, actuators, commands, observations, rewards, terminations: pure JAX, shared by simulator and imagination
  sim.py         MuJoCo Warp simulator: scene, physics, contacts, resets, pushes, forward kinematics with velocities
  train.py       the learner: collection, motion priors, world model, BPTT (+ critic, PPO), checkpoints; Config = all flags
  evaluate.py    evaluation of checkpoints
  amp.py         AMP discriminator (WGAN-GP), frame features, style reward
  getup_amp.py   demo windows of the get-up prior
  nets.py        world model and policy networks, ODE actor, running normalizers
  critic.py      value function and its TD(lambda) fit
  buffer.py      device replay buffer
  resets.py      reset banks from motion clips
  optim.py       MuonAdam
  kin.py, rot.py forward kinematics and rotations
assets/g1_mjlab.xml   the G1 MJCF (visual meshes: tools/download_meshes.sh)
tools/         download_data.sh, convert_motions.py, download_meshes.sh, render.py, check_heads.py
tests/         CPU tests: actuator and kinematics against MuJoCo, shapes

Acknowledgements and citation

InstantAMP stands on the following works. If you use it, please cite them.

• SuperMimic (Liao et al., 2026): the method, a differentiable world model with analytic policy gradients through imagined rollouts. Paper · Code
• InstantTrack by Huy-Quang Dao: G1 motion tracking in minutes on one GPU, the origin of much of the world-model training code (through BlinkMimic and XMimic). Code
• XSoma (Bohlinger and Peters, 2026): the JAX and MuJoCo Warp implementation that InstantAMP is built on. Paper · Code
• ODE Policy (Jia et al., 2026): the ODE actor (--ode 1). Paper · Code

Also thanks to:

• AMP (Peng et al., 2021, arXiv:2104.02180);
• mjlab, for the G1 model and actuator constants;
• the code lineage of XMimic and supermimic_vibed;
• legged_lab and LAFAN1, for the motion data;
• MuJoCo, MuJoCo Warp, NVIDIA Warp and JAX.

InstantAMP is an independent project. It is not affiliated with or endorsed by the authors above.

@article{liao2026supermimic,
  title   = {SuperMimic: Scaling Humanoid Whole-Body Control with One Differentiable World Model},
  author  = {Liao, Qiayuan and Shao, Yiyang and Long, Junfeng and Zakka, Kevin and Horowitz, Roberto and Abbeel, Pieter and Sreenath, Koushil},
  journal = {arXiv preprint arXiv:2610.28480},
  year    = {2026}
}

@misc{dao2026instanttrack,
  title        = {InstantTrack: Whole-Body Motion Tracking for the Unitree G1 in Minutes on One GPU},
  author       = {Dao, Huy-Quang},
  year         = {2026},
  howpublished = {\url{https://github.com/Huy-Quang-Dao/instantTrack}}
}

@misc{bohlinger2026xmimicxsoma,
  title  = {XMimic and XSoma: Learning Humanoid Motion Tracking and Locomotion in Under a Minute},
  author = {Bohlinger, Nico and Peters, Jan},
  year   = {2026},
  url    = {https://nico-bohlinger.github.io/papers/2026_xmimic_xsoma.pdf}
}

@article{jia2026odepolicy,
  title   = {ODE Policy: Robot Learning via Action Dynamics, Not Actions},
  author  = {Jia, Jindou and Li, Jingliang and Zuo, Kuangji and An, Tuo and Hu, Yuxuan and Zhao, Mengfei and Lyu, Bofan and Bai, Jiaqi and Zhu, Daoyuan and Dai, Hengyu and Zhang, Junpeng and Yang, Jianfei},
  journal = {arXiv preprint arXiv:2611.01729},
  year    = {2026}
}

@article{peng2021amp,
  title   = {AMP: Adversarial Motion Priors for Stylized Physics-Based Character Control},
  author  = {Peng, Xue Bin and Ma, Ze and Abbeel, Pieter and Levine, Sergey and Kanazawa, Angjoo},
  journal = {ACM Transactions on Graphics},
  volume  = {40},
  number  = {4},
  year    = {2021}
}

License

InstantAMP is released under the MIT License. It derives from XSoma (MIT), whose copyright notice is kept. The bundled G1 MJCF comes from mjlab (Apache-2.0) and derives from Unitree's BSD-3-Clause description. The motion clips are not included; their own licenses apply, and LAFAN1 is non-commercial. Details: THIRD_PARTY_NOTICES.md.