RL

A G1 tracks a LAFAN1 dance.	A Go1 follows joystick commands.

robokit.helpers.rl trains policies on batched GPU simulations, with features including PPO, CUDA-graphed updates, and coarse-to-fine physics.

This tutorial covers the main RL workflows. The examples' envs run on MuJoCo Warp (MjWarpSim); install it with the mj and mjwarp extras, and PyTorch with a pt-* extra. Complete runnable examples are in src/robokit/examples/rl.

How RL Works

Each PPO update follows three steps:

1. Roll the actor out for horizon steps in every world of the env.
2. Update the actor and the critic on that rollout over several epochs.
3. Adapt the learning rate to the policy's KL divergence.

The env's tensors stay on the GPU, so no step waits for the host. After a few warm-up updates, each update replays as one CUDA graph.

Environments

An env implements Env: n worlds, the actor's and the critic's observation sizes (obs_dims), the action bound, and the names of its reward terms. Its methods take and return GPU tensors:

The examples wrap MjWarpSim worlds and read sim.data.torch.qpos, qvel and ctrl, zero-copy torch views.

Basic PPO

A PPO holds an actor and a critic. Each update(env) rolls out once and updates once; the caller owns the loop. This trains on a toy env where actions near zero pay:

from robokit.helpers.rl import ppo

learner = ppo.PPO(env.obs_dims, len(env.bound), ppo.PPOConfig())
for _ in range(10):
    learner.update(env)

score = ppo.evaluate(env, learner.actor, steps=100)  # each reward term's mean per step, and terminations

from types import SimpleNamespace

import torch

obs, done = torch.zeros(64, 3, device="cuda"), torch.zeros(64, dtype=torch.bool, device="cuda")
env = SimpleNamespace(
    n=64,
    obs_dims=(3, 3),
    bound=torch.ones(2, device="cuda"),
    names=["effort"],
    obs=lambda: (obs, obs),
    step=lambda a: (-(a**2).sum(1), done, done, (a**2).mean()[None]),
    reset=id,
    curriculum=id,
)
from robokit.helpers.rl import ppo

learner = ppo.PPO(env.obs_dims, len(env.bound), ppo.PPOConfig())
for _ in range(10):
    learner.update(env)

score = ppo.evaluate(env, learner.actor, steps=100)  # each reward term's mean per step, and terminations
assert set(score) == {"effort", "terminations"}

learner.actor(obs) returns the mean action. PPOConfig holds the PPO settings, such as horizon, epochs, minibatches, lr, kl, entropy, and hidden.

Coarse-to-Fine Training

A policy can train on cheaper physics first. Build the coarse env's scene with MjWarpSim.build_coarse_spec and a CoarseConfig, then let train mix it in: over the first half of the time, an update draws it with a chance falling to 0.

actor = learner.train(fine, minutes, coarse_env=coarse, coarse=0.5)

Examples

Run one with uv run src/robokit/examples/rl/13_cartpole.py swingup --minutes 1.