PPO

StableTested and in active use. Breaking changes will be called out.

PPO is an on-policy policy-gradient method. It improves the policy in small, clipped steps to stay stable: collect a fresh rollout, estimate advantages, run several update epochs over that data, then throw it away.

It needs more environment steps than SAC and is considerably harder to destabilise. With a batched GPU environment, the sample cost is usually the cheaper trade.

LunarLander — PPO8 obs · 4 discrete actions

Papers

• Proximal Policy Optimization AlgorithmsSchulman et al. · 2017 · arXiv:1707.06347
• High-Dimensional Continuous Control Using Generalized Advantage EstimationSchulman et al. · 2015 · ICLR 2016 · arXiv:1506.02438GAE — the advantage estimator.

The learning loop

environment interactionpolicy updatevalue / advantage

clipped surrogateL = E[ min( r·Â, clip(r, 1±ε)·Â ) ] where r = π / π_old

The structural difference from SAC is the buffer. SAC’s replay persists for a million transitions and is sampled randomly. PPO’s rollout buffer is filled, consumed by N_EPOCHS passes, and discarded. Nothing survives the iteration.

The clipped surrogate. The update maximises the ratio r = π(a|s) / π_old(a|s) weighted by the advantage — but clipped to [1−ε, 1+ε]. Once the new policy has moved far enough from the old one, the gradient flattens, so no single update can take a destructive step. That clip is essentially the whole algorithm.

GAE advantages. Generalized Advantage Estimation trades bias against variance with a parameter λ. At λ=1 it is the raw Monte-Carlo return (unbiased, noisy); at λ=0 it is the one-step TD error (biased, stable). 0.95 is the usual compromise.

Multiple epochs. Because the clip bounds how far the policy can move, the same rollout can safely be reused for several gradient passes — which is what makes an on-policy method affordable at all.

Building it in noeira

From examples/half_cheetah/ppo_half_cheetah_training.mojo:

from noeira.nn.constants import DT
from noeira.nn.combinators.sequential import Sequential
from noeira.nn.primitives.linear import Linear
from noeira.nn.primitives.activations import Tanh
from noeira.deep_agents.ppo import PPOAgent
from noeira.envs.half_cheetah import HalfCheetah, HalfCheetahConfig

comptime OBS_DIM = HalfCheetahConfig.OBS_DIM      # 17
comptime ACT_DIM = HalfCheetahConfig.ACTION_DIM   #  6
comptime HIDDEN = 256
comptime ROLLOUT_LEN = 2_048
comptime MINIBATCH = 64
comptime N_EPOCHS = 10
comptime LOG_STD_INIT: Scalar[DT] = -0.5
comptime MAX_TORQUE: Scalar[DT] = 1.0

comptime CriticNet = Sequential[
    Linear[OBS_DIM, HIDDEN], Tanh[HIDDEN],
    Linear[HIDDEN, HIDDEN], Tanh[HIDDEN],
    Linear[HIDDEN, 1],
]


def main() raises:
    var agent = PPOAgent[
        "cpu",
        ActorNet,
        CriticNet,
        OBS_DIM,
        ACT_DIM,
        ROLLOUT_LEN,
        MINIBATCH,
        N_EPOCHS,
    ](
        actor_lr=3e-4,
        critic_lr=1e-3,
        gamma=0.99,
        gae_lambda=0.95,
        clip_eps=0.2,
        entropy_coef=0.0,
        action_scale=MAX_TORQUE,
        log_std_init=LOG_STD_INIT,
        max_grad_norm=0.5,
    )

    # CleanRL-style log-std init. The trainer leaves this to the caller because
    # Mojo nightly cannot reflect into Sequential's variadic children
    # generically. Index 4 is the GaussianHead.
    agent.trainer.actor.children[4].set_log_std_init["cpu"](LOG_STD_INIT)

    var env = HalfCheetah[DT, TERMINATE_ON_UNHEALTHY=False]()
    _ = agent.train_single[HalfCheetah[DT, TERMINATE_ON_UNHEALTHY=False]](
        env, 1_000_000, print_every=50_000, verbose=True,
    )

Note Tanh activations, not ReLU. That is the CleanRL convention for PPO on MuJoCo tasks and it matters — ReLU trunks are noticeably less stable here.

Compile-time parameters

PPOAgent[train_target, ACTOR, CRITIC, OBS_DIM, ACT_DIM,
         ROLLOUT_LEN, MINIBATCH, N_EPOCHS, N_ENVS = 1]

Unlike the off-policy agents, PPO takes its dimensions directly rather than through a sample block — there is no replay buffer to carry them. ROLLOUT_LEN, MINIBATCH and N_EPOCHS are compile-time because they determine buffer layout.

Key runtime arguments

Argument	Typical	Notes
clip_eps	0.2	the surrogate clip — the core knob
gae_lambda	0.95	bias/variance trade
gamma	0.99	discount
entropy_coef	0.0–0.01	exploration bonus; 0.0 for MuJoCo
max_grad_norm	0.5	canonical PPO gradient clip
log_std_init	-0.5	initial policy standard deviation

max_grad_norm=0.5 is not optional

The canonical PPO gradient-norm clip closed a real gap here — on Pendulum at 200k steps it moved the result from −230 to −135 at the same seed. It is in the original paper and it earns its place.

Batched GPU training

For the batched driver, set N_ENVS and use train:

comptime N_ENVS = 32
comptime EnvT = Phyics3dBatchedEnv[HalfCheetahModel, HalfCheetahConfig, N_ENVS]

var agent = PPOAgent[
    "gpu", ActorNet, CriticNet, OBS_DIM, ACT_DIM,
    ROLLOUT_LEN, MINIBATCH, N_EPOCHS, N_ENVS,
](ctx=ctx, ...)

This is where PPO becomes practical: its appetite for samples is satisfied by stepping 32 environments per iteration on-device.

Keeping the rollout on the device

By default the batched driver stages the rollout on the host. Every env step it copies the observation, reward and done flags back, samples on the host and records into host lists. Each minibatch then uploads its tensors and reads its losses back. Three compile-time flags on train (and on PPODiscreteAgent’s train_batched) remove those round trips. All three are off by default:

Flag	What moves to the device
USE_TRAIN_CUDA_GRAPH	the K-epoch update: one minibatch step is captured and replayed per minibatch
DEVICE_ROLLOUT	action sampling, recording and GAE; episode returns come back through a ring every 32 steps
USE_ENV_CUDA_GRAPH	the env step and the reset, captured (needs DEVICE_ROLLOUT and an env step that draws no random numbers)

The captured update is bit-identical to the eager one. The device rollout samples with a device RNG, so its trajectories differ from the host rollout’s, but with and without its graphs it writes byte-identical checkpoints (tests/nn/test_ppo_train_graph_parity.mojo, discrete and continuous).

LunarLander on an RTX 5090: 11.8k env steps/s staged on the host, 17.2k with the captured update, and 18.8k with everything on the device and captured. With the cuBLAS GEMM routing on top it reaches 25.4k.

pixi run -e nvidia mojo build -I . -D PPO_GPU -D PPO_DEVICE_ROLLOUT -D PPO_ENV_GRAPH -D PPO_TRAIN_GRAPH \
    examples/lunar_lander/ppo_lunar_lander_time_to_solve.mojo -o ppo_ll_gpu_dev
pixi run ./ppo_ll_gpu_dev 1 ppo_ll_cuda_dev_s1.ckpt

Run the binary through pixi run: the graphs need the CUDA interceptor that the nvidia environment preloads. Without it they disable themselves and the same device path runs eagerly.

run_ppo_vec: any GPU-batched environment

run_ppo_vec (noeira/deep_agents/training/ppo_vec_driver.mojo) is continuous PPO over any GPU-batched environment, with the rollout always on the device. The env step, the reset and the update can each be captured (ENV_GRAPH, TRAIN_GRAPH). Its options are opt-in, most of them fields of PPOVecConfig:

• Observation. The env’s row followed by the last HIST executed actions. FRAMES stacks a per-lane history of K frames, newest first, as IsaacLab does. E_ACTOR splits each frame into actor words and privileged words: an actor with a Slice at its head reads only the actor prefix, which gives a privileged critic with no trainer change.
• Normalisation. Running observation statistics and CleanRL’s NormalizeReward, each switchable.
• Schedules. The learning rate and the entropy coefficient annealed linearly or held fixed, or rsl_rl’s adaptive-KL learning rate (adaptive_kl). Adaptive KL is applied once per update, so the update stays one graph.
• Penalties. Action rate w·|a_t − a_{t−1}|² (act_rate_w) and smoothness w·|a_t − 2a_{t−1} + a_{t−2}|² (act_smooth_w).
• Mirror symmetry. An opt-in loss on the actor mean (PPOActorLoss.enable_mirror, noeira/deep_agents/ppo/mirror.mojo).
• A divergence guard. A lane whose observation or reward goes non-finite or out of bounds is ended as terminated and counted.

A checkpoint is the trainer’s v3 checkpoint plus obs_norm.txt. On Hopper (examples/hopper/ppo_hopper_vec_gpu.mojo, CleanRL’s MuJoCo recipe, 128 lanes), the last-100 return at 2M steps is 1639 / 2911 / 2722 over seeds 1–3, at about 7k env steps/s. The G1 walker trains through it.

Discrete actions

PPO handles both action types, through separate packages:

from noeira.deep_agents.ppo import PPOAgent               # continuous, Gaussian head
from noeira.deep_agents.ppo_discrete import PPODiscreteAgent  # categorical head

PPODiscreteAgent has a train_batched entry point for GPU-batched discrete environments — the arcade games and Procgen.

Watch out for

Time limits break GAE

The classic PPO bug, and Pendulum is where it bites: the environment has no terminal state, so every episode ends by truncation. Treating that as termination truncates the value bootstrap and corrupts every advantage estimate in the rollout.

Use was_terminated() rather than done when computing GAE. See Traits.

The on-policy drivers store each row’s pre-reset next observation. A truncated row bootstraps from the value of its own final state, and the advantage recursion stops at the episode boundary. Until October 2026 they bootstrapped a truncation from the next episode’s first state, and every task-family run was affected. With the fix, lift_brick’s greedy success went from 97.9 / 92.7 / 92.4 % to 96.3 / 95.5 / 94.7 % over three seeds. tests/nn/test_ppo_gae_truncation.mojo gates the host and device paths against the formula written out.

Observation normalization. PPO is more sensitive to observation scale than off-policy methods, because the advantage estimates feed directly into the policy gradient. If a batched environment wrapper is applying normalization, make sure it is applied consistently between rollout collection and update.

A2C

A2C is PPO with a single epoch and no clipping — a degenerate case, and implemented as such here:

from noeira.deep_agents.a2c import A2CAgent

Simpler, but strictly worse in practice. Use PPO unless you specifically want the comparison.

Where to run it

pixi run mojo run -I . examples/cartpole/cartpole_ppo.mojo
pixi run mojo run -I . examples/half_cheetah/ppo_half_cheetah_training.mojo
pixi run -e nvidia mojo run -I . examples/craftax_classic/ppo_training_gpu.mojo

# run_ppo_vec on Hopper, env step and update captured
pixi run -e nvidia mojo build -I . -D PPO_VEC_ENV_GRAPH -D PPO_VEC_TRAIN_GRAPH \
    examples/hopper/ppo_hopper_vec_gpu.mojo -o ppo_hopper_vec
pixi run -e nvidia ./ppo_hopper_vec --steps 2000000 --seed 1

Next

• SAC — off-policy, for comparison
• MuZero — planning instead of pure policy gradient
• Algorithms overview — the full catalogue