Training

Two scripts cover the clone and reinforce stages. examples/bc_ppo_warp.py is the default recommendation: it BC-pretrains the actor on teleoperation demonstrations, then runs GPU-parallel PPO with a persistent BC loss anchoring the actor to those demos. examples/ppo_warp.py is the same recipe with no demos, trained from scratch, and --no-use-demos makes bc_ppo_warp.py recover it exactly. examples/ppo.py is a slow single-environment MuJoCo reference, useful for sanity checks but not the recommended path.

Quick start

The defaults target WarpPickLift-v1 with the published johnsutor/MuJoCoPickLift-v1 demonstrations, so one command runs the whole thing:

uv run --extra warp --extra train python examples/bc_ppo_warp.py

Watch it train:

tensorboard --logdir runs

Evaluate the saved checkpoint deterministically, with no exploration noise:

uv run --extra warp python examples/eval_warp.py \
    --env-id WarpPickLift-v1 \
    --checkpoint "runs/WarpPickLift-v1__*/best_agent.pt"

examples/ppo_warp.py can also render periodic transfer videos in the matching MuJoCo* environment with --capture-video.

Results

Measured on one RTX 5090. Success is the recent completed-episode success rate reported by the Warp training rollout at the listed step budget. The two scripts are reported separately because they are different recipes: a task with no vanilla baseline is not a broken task, and a task solved only under demo seeding is still solved.

Under vanilla PPO, WarpPickLift-v1 seed 5 stalls at a grasp-and-hold-at-the-table local optimum with best_success=0.037; demo seeding alone fixes it at the same 30M budget, which is the whole argument for the clone stage. The WarpPickAndPlace-v1 seed-1 checkpoint also transfers across backends: a 100-episode evaluation in MuJoCo reports success_rate=0.870.

Recommended commands

WarpTouch-v1, WarpLookAt-v1, and WarpMove-v1 need nothing beyond a step budget:

uv run --extra warp --extra train python examples/ppo_warp.py \
    --env-id WarpTouch-v1 \
    --total-timesteps 5000000

WarpPickLift-v1 under vanilla PPO. This repeats the tuned defaults explicitly so copied runs keep the seed-validated recipe even if script defaults change later:

uv run --extra warp --extra train python examples/ppo_warp.py \
    --env-id WarpPickLift-v1 \
    --total-timesteps 30000000 \
    --num-minibatches 32 \
    --update-epochs 10 \
    --ent-coef 0.03 \
    --ent-coef-final 0.005 \
    --max-grad-norm 0.5 \
    --target-kl None

WarpPickAndPlace-v1, the validated demo-seeded run. The success bonus, the long budget, and the decoupled anneal horizon with a learning-rate floor are all load-bearing here:

uv run --extra warp --extra train python examples/bc_ppo_warp.py \
    --env-id WarpPickAndPlace-v1 \
    --demo-repo johnsutor/MuJoCoPickAndPlace-v1 \
    --success-bonus 50 \
    --total-timesteps 160000000 \
    --anneal-timesteps 80000000 \
    --lr-min-frac 0.1

If a pick task stalls at low return with 0% success, check the entropy and optimizer settings against the commands above before tuning anything else. Lift discovery on GPU is stochastic.

Behavior cloning

The clone stage is the BC-pretrain phase inside bc_ppo_warp.py. It regresses the actor mean onto demo actions before any online rollout, so PPO starts near successful behavior instead of a random init. There is no separate BC step to run.

Point it at your own recording with --demo-repo, which takes a Hugging Face repo id or a local dataset path. Every field of the script's Args dataclass is exposed on the command line through tyro:

uv run --extra warp --extra train python examples/bc_ppo_warp.py \
    --demo-repo your-user/MuJoCoPickLift-v1 \
    --env-id WarpPickLift-v1

To judge clone quality, watch pretrain/bc_loss in TensorBoard during the pretrain phase. A low, flat curve means the actor has matched the demo actions. --bc-coef is what makes fine-tuning reinforce the clone rather than forget it; --bc-anneal-steps fades that anchor if you want the policy to drift toward its own discoveries.

Baseline hyperparameters

The recipe uses fixed-horizon episodes, staggered resets, observation normalization, return-based reward scaling, and entropy annealing. Treat the annealing horizon as a real hyperparameter: shortening --total-timesteps also shortens the learning-rate and entropy schedules, unless you decouple it with --anneal-timesteps. The pick tasks use a strong entropy warm-start with a nonzero floor because the GPU Warp contact path is not bitwise deterministic.

Caveats

Colab

bc_ppo_warp_colab.ipynb runs the demo-seeded pipeline on WarpPickLift-v1: it downloads the published demos, BC-pretrains the actor, fine-tunes with PPO, and shows a rollout video.

ppo_warp_colab.ipynb runs vanilla PPO with the same ENV_ID and step-budget choices, embedded TensorBoard, evaluation, and a rollout video.