Training with PPO

SO101-Nexus uses examples/ppo_warp.py as the default PPO baseline for RL. It is a CleanRL-style MLP actor-critic trainer for the GPU-batched Warp*-v1 environments. Rollouts stay on the CUDA device, so this is the path to use for serious PPO experiments.

examples/ppo.py is kept as a slower MuJoCo reference script for single-machine sanity checks. Prefer the Warp script for baseline numbers, Colab runs, and new tuning work.

Per-environment starting points

These are grounded baselines from one RTX 5090 using the common settings below. "Success" is the recent completed-episode success rate reported by the Warp training rollout. PickLift reports seed-validated results from seeds 1, 2, 3, 4, and 5.

Recommended commands. The PickLift command 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 WarpTouch-v1 \
    --total-timesteps 5000000

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

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

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

If PickLift stalls at low return with 0% success, first verify that you are using the recommended entropy and optimizer settings above. Lift discovery is still stochastic on GPU, but the default recipe was validated across seeds 1, 2, 3, 4, and 5 (4 of 5 solved).

Quick start

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

Open TensorBoard while training:

tensorboard --logdir runs

Baseline recipe

The baseline 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. WarpPickLift-v1 uses a strong entropy warm-start (ent_coef=0.03) with a nonzero floor (ent_coef_final=0.005) and a CleanRL-style update budget because the GPU Warp contact path is not bitwise deterministic.

Common settings for the solved baselines:

Evaluation

Evaluate a saved checkpoint deterministically with the matching Warp environment:

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* backend with --capture-video.

Demo-seeded PPO

examples/bc_ppo_warp.py is the same PPO recipe plus behavior-cloning (BC) seeding from teleoperation demonstrations: the actor is BC-pretrained on the demos before online PPO starts, and a persistent BC loss (--bc-coef, default 0.1) keeps anchoring it toward demo actions throughout training. --use-demos false recovers ppo_warp.py exactly.

Run it on your own dataset, or the published johnsutor/MuJoCoPickLift-v1 default:

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

This targets the one seed-fragility gap in the PPO baseline above: seed 5 gets stuck at a grasp-hold-at-table local optimum under vanilla PPO (best_success=0.037). Same 30M-step recipe, demo-seeding alone rescues it to best_success=0.993, final_success=0.983. For running BC on your own dataset, inspecting the clone, and the action-space details, see Behavior Cloning. Full design notes and results are in examples/README.md.

Colab

For browser training, open examples/ppo_warp_colab.ipynb. The notebook exposes the same ENV_ID and step-budget choices, embeds TensorBoard, trains, evaluates, and displays a rollout video.

For the demo-seeded (BC + PPO) pipeline on WarpPickLift-v1, open examples/bc_ppo_warp_colab.ipynb. It downloads the published johnsutor/MuJoCoPickLift-v1 demos, BC-pretrains the actor, fine-tunes with PPO, and shows a rollout video.