Module tutorial_02_cartpole_ppo 

Tutorial 2 — Solving CartPole with PPO.

See the rendered Markdown at docs/tutorials/02-cartpole-ppo.md.

Tutorial 2 — Solving CartPole with PPO

You will: train a PPO agent to balance CartPole, this time letting PPOTrainerBurn own the update step, and learn to read the learning curve it produces. Prerequisites: Tutorial 1 — you should recognize the rollout → advantage → update loop and Burn’s move-through optimizer. Time: ~15 minutes.

Tutorial 1 hand-wrote the whole PPO update to demystify it. Now we graduate to a real sequential task and a real trainer. Two things change from the bandit:

1. The future matters. In CartPole, the action you take now affects the states you’ll see for the rest of the episode. That means we need proper temporal credit assignment — GAE (Generalized Advantage Estimation) — instead of the bandit’s one-line reward − value.
2. We stop writing the loss by hand. PPOTrainerBurn owns the clipped surrogate, value loss, entropy bonus, gradient clipping, and the optimizer step. We keep control of the rollout and the advantage computation, and hand the trainer a batch of tensors. This is the division of labor you’ll use for every on-policy trainer in Thrust.

The problem: CartPole-v1

A pole is hinged on a cart; you push the cart left (0) or right (1) each step to keep the pole upright. The observation is four floats ([cart position, cart velocity, pole angle, pole angular velocity]), the reward is +1 per step the pole stays up, and the episode ends when the pole falls or at 500 steps. So episode reward == episode length, and a solved agent scores near 500. Random pushing scores ~22.

GAE, briefly

The advantage answers “how much better than expected was this action?”. For a sequential task, “expected” has to account for the future, so GAE walks the rollout backwards, discounting future rewards by gamma and blending multi-step estimates with gae_lambda:

delta_t = r_t + gamma * V(s_{t+1}) - V(s_t)
A_t     = delta_t + gamma * lambda * A_{t+1}

gamma = 0.99 says “future reward is almost as valuable as immediate reward” (appropriate when episodes are long and every step matters). gae_lambda = 0.95 trades a little bias for much lower variance. Contrast this with Tutorial 1, where we set both to 0 precisely because the bandit has no future to bootstrap from. We compute GAE ourselves so you can see it; the helper is at the bottom of the code block.

The trainer surface

PPOConfig is where the on-policy knobs live. The important ones:

• learning_rate, n_epochs, batch_size — optimization.
• gamma, gae_lambda — credit assignment (must match how you compute advantages).
• clip_range — the PPO trust-region clip (the same 0.2 from Tutorial 1).
• vf_coef, ent_coef — value-loss and entropy weights.
• max_grad_norm — gradient clipping for stability.

You build a PPOConfig, wrap your policy and a BurnOptimizer in a PPOTrainerBurn, then call train_step once per rollout. See PPO_BEST_PRACTICES.md for how to tune these.

The code

This runs a short training loop (small budget so it’s fast in CI; bump TOTAL_TIMESTEPS for a real run). It is a doc-test, so it always compiles against the current API.

use burn::{
    backend::{Autodiff, NdArray},
    optim::AdamConfig,
    tensor::{Int, Tensor, TensorData},
};
use thrust_rl::prelude::*;
use thrust_rl::train::optimizer::BurnOptimizer;

type Backend = Autodiff<NdArray<f32>>;

// --- Hyperparameters -------------------------------------------------------
const NUM_ENVS: usize = 16; // parallel CartPole copies
const NUM_STEPS: usize = 256; // rollout length per env
const TOTAL_TIMESTEPS: usize = 8_192; // tiny for CI; use ~200_000 for real
const LEARNING_RATE: f64 = 3e-4;
const HIDDEN_DIM: usize = 128;
const GAMMA: f32 = 0.99;
const GAE_LAMBDA: f32 = 0.95;
const SEED: u64 = 0; // seed policy init + resets for reproducibility

let device = Default::default();

// --- Environment -----------------------------------------------------------
let probe = CartPole::new();
let obs_dim = probe.observation_space().shape[0];
let action_dim = match probe.action_space().space_type {
    SpaceType::Discrete(n) => n,
    SpaceType::Box => panic!("CartPole is discrete"),
};
let mut env_pool = EnvPool::new(CartPole::new, NUM_ENVS);

// --- Policy: a seeded 2-layer ReLU MLP with orthogonal init ----------------
// Seeding makes the run reproducible, which is what lets a learning-curve CSV
// be overlaid across algorithms on the same axis (see below).
let policy_config = MlpBurnConfig {
    num_layers: 2,
    hidden_dim: HIDDEN_DIM,
    use_orthogonal_init: true,
    activation: BurnActivation::ReLU,
    seed: Some(SEED),
};
let policy = MlpBurnPolicy::<Backend>::with_config(obs_dim, action_dim, policy_config, &device);

// --- Optimizer + trainer ---------------------------------------------------
let inner_opt = AdamConfig::new().init();
let burn_opt: BurnOptimizer<Backend, MlpBurnPolicy<Backend>, _> =
    BurnOptimizer::new(inner_opt, LEARNING_RATE);

let ppo_config = PPOConfig::new()
    .learning_rate(LEARNING_RATE)
    .n_epochs(10)
    .batch_size(128)
    .gamma(GAMMA as f64)
    .gae_lambda(GAE_LAMBDA as f64)
    .clip_range(0.2)
    .clip_range_vf(0.2)
    .vf_coef(0.5)
    .ent_coef(0.01)
    .max_grad_norm(0.5)
    // Disable KL early-stop for a short, small-budget run.
    .target_kl(1.0);

let mut trainer = PPOTrainerBurn::new(ppo_config, policy, burn_opt).unwrap();

let cap = NUM_STEPS * NUM_ENVS;
let num_updates = TOTAL_TIMESTEPS / cap;

// Rollout buffers, reused each update.
let mut buf_obs: Vec<f32> = Vec::with_capacity(cap * obs_dim);
let mut buf_actions: Vec<i64> = Vec::with_capacity(cap);
let mut buf_log_probs: Vec<f32> = Vec::with_capacity(cap);
let mut buf_values: Vec<f32> = Vec::with_capacity(cap);
let mut buf_rewards: Vec<f32> = Vec::with_capacity(cap);
let mut buf_dones: Vec<f32> = Vec::with_capacity(cap);

let mut observations = env_pool.reset();
let mut episode_lengths = [0u32; NUM_ENVS];
let mut completed: Vec<u32> = Vec::new();

for _update in 0..num_updates {
    buf_obs.clear();
    buf_actions.clear();
    buf_log_probs.clear();
    buf_values.clear();
    buf_rewards.clear();
    buf_dones.clear();

    // --- Phase 1: rollout -------------------------------------------------
    for _step in 0..NUM_STEPS {
        let obs_flat: Vec<f32> = observations.iter().flatten().copied().collect();
        let obs_t: Tensor<Backend, 2> =
            Tensor::from_data(TensorData::new(obs_flat, [NUM_ENVS, obs_dim]), &device);
        let (actions, log_probs, values) = trainer.policy().get_action_host(obs_t);
        let results = env_pool.step(&actions);

        for env_id in 0..NUM_ENVS {
            buf_obs.extend_from_slice(&observations[env_id]);
            buf_actions.push(actions[env_id]);
            buf_log_probs.push(log_probs[env_id]);
            buf_values.push(values[env_id]);
            buf_rewards.push(results[env_id].reward);

            let done = results[env_id].terminated || results[env_id].truncated;
            buf_dones.push(if done { 1.0 } else { 0.0 });

            episode_lengths[env_id] += 1;
            observations[env_id] = results[env_id].observation.clone();
            if done {
                completed.push(episode_lengths[env_id]);
                trainer.increment_episodes(1);
                episode_lengths[env_id] = 0;
                observations[env_id] = env_pool.reset_env(env_id).unwrap();
            }
        }
    }

    // --- Phase 2: GAE (needs a value bootstrap for the final observation) --
    let last_obs_flat: Vec<f32> = observations.iter().flatten().copied().collect();
    let last_obs_t: Tensor<Backend, 2> =
        Tensor::from_data(TensorData::new(last_obs_flat, [NUM_ENVS, obs_dim]), &device);
    let (_, _, last_values) = trainer.policy().get_action_host(last_obs_t);

    let (advantages, returns) = compute_gae(
        &buf_rewards, &buf_values, &buf_dones, &last_values, GAMMA, GAE_LAMBDA, NUM_STEPS, NUM_ENVS,
    );

    // --- Phase 3: hand the batch to the trainer ---------------------------
    let batch = cap;
    let obs_b: Tensor<Backend, 2> =
        Tensor::from_data(TensorData::new(buf_obs.clone(), [batch, obs_dim]), &device);
    let actions_b: Tensor<Backend, 1, Int> =
        Tensor::from_data(TensorData::new(buf_actions.clone(), [batch]), &device);
    let old_log_probs_b: Tensor<Backend, 1> =
        Tensor::from_data(TensorData::new(buf_log_probs.clone(), [batch]), &device);
    let old_values_b: Tensor<Backend, 1> =
        Tensor::from_data(TensorData::new(buf_values.clone(), [batch]), &device);
    let advantages_b: Tensor<Backend, 1> =
        Tensor::from_data(TensorData::new(advantages, [batch]), &device);
    let returns_b: Tensor<Backend, 1> =
        Tensor::from_data(TensorData::new(returns, [batch]), &device);

    // The trainer runs n_epochs of minibatched clipped-surrogate updates and
    // returns per-update stats. The closure tells it how to score
    // (obs, action) batches with the current policy.
    let _stats = trainer
        .train_step(
            obs_b,
            actions_b,
            old_log_probs_b,
            old_values_b,
            advantages_b,
            returns_b,
            |p, o, a| p.evaluate_actions(o, a),
        )
        .unwrap();
}

// A short CI run won't fully solve CartPole, but the trainer should have run.
// We assert only that the loop executed, to keep the doc-test fast and
// deterministic.
assert!(num_updates >= 1);
let _ = completed;

/// Per-env GAE. `rewards`, `values`, `dones` are flat `[T * N]` row-major
/// (step-major: index = step * num_envs + env_id). `last_values[n]` is the
/// value bootstrap for env `n` at the step just past the end of the rollout.
#[allow(clippy::too_many_arguments)]
fn compute_gae(
    rewards: &[f32],
    values: &[f32],
    dones: &[f32],
    last_values: &[f32],
    gamma: f32,
    gae_lambda: f32,
    num_steps: usize,
    num_envs: usize,
) -> (Vec<f32>, Vec<f32>) {
    let cap = num_steps * num_envs;
    let mut advantages = vec![0.0_f32; cap];
    let mut returns = vec![0.0_f32; cap];
    let mut last_gae = vec![0.0_f32; num_envs];
    // Walk the rollout in reverse per env so each advantage can fold in the
    // (already-computed) advantage of the following step.
    for t in (0..num_steps).rev() {
        for n in 0..num_envs {
            let idx = t * num_envs + n;
            let next_value = if t == num_steps - 1 {
                last_values[n]
            } else {
                values[(t + 1) * num_envs + n]
            };
            // If the episode ended at this step, don't bootstrap across the
            // boundary — the next state belongs to a fresh episode.
            let next_nonterminal = 1.0 - dones[idx];
            let delta = rewards[idx] + gamma * next_value * next_nonterminal - values[idx];
            last_gae[n] = delta + gamma * gae_lambda * next_nonterminal * last_gae[n];
            advantages[idx] = last_gae[n];
            returns[idx] = advantages[idx] + values[idx];
        }
    }
    (advantages, returns)
}

Reading the learning curve

The packaged train_cartpole_modern example is the same loop with a bigger budget and one extra feature: set CURVE_CSV and it writes one env_steps,mean_episode_reward row per update.

TOTAL_TIMESTEPS=200000 CURVE_CSV=/tmp/ppo.csv \
    cargo run --release --features training --example train_cartpole_modern

Because CartPole’s reward is +1/step, mean_episode_reward is the mean episode length. Plot env_steps (x) against mean_episode_reward (y) and you should see it climb off the ~22 random baseline and head toward 500. What to look for:

• A rising curve = credit assignment is working.
• A curve stuck near ~22 = something is off (a common cause is a gamma/gae_lambda mismatch between PPOConfig and your compute_gae call — they must agree).
• A curve that climbs then crashes = the policy collapsed; lower the learning rate or the entropy coefficient, or re-enable target_kl.

The seed makes runs reproducible, so two algorithms (say PPO here and A2C from the train_cartpole_a2c example) can be overlaid on the same env, seed, and budget for an honest comparison.

Try it yourself

• Mismatch the discount on purpose: pass gamma = 0.0 to compute_gae while leaving PPOConfig::gamma(0.99) — the returns the critic is trained on no longer match the advantages, and learning stalls. Concrete proof the two must agree.
• Shrink the pool: set NUM_ENVS = 1. Fewer parallel envs means noisier gradients per update; you’ll need more updates to reach the same bar.
• Full run: drop the CI-sized TOTAL_TIMESTEPS and run the packaged example above with the CSV, then plot it.

Next

The path splits from here — see the tutorial index for the full dependency-ordered series. Next up (Tutorial 3) is off-policy training with DQN: replay buffers, target networks, and when to reach for DQN over PPO.