Expand description
Tutorial 5 — Memory and POMDPs with recurrent PPO.
See the rendered Markdown at
docs/tutorials/05-memory-pomdps.md.
§Tutorial 5 — Memory and POMDPs with recurrent PPO
You will: train a recurrent (LSTM) PPO agent on
FlickeringCartPole— a version of CartPole whose observation is randomly blanked to zeros — and watch it beat a memoryless MLP baseline that sees the identical stream. Along the way you’ll learn what a POMDP is, when memory is worth its cost, and how a recurrent rollout differs from the feedforward one from Tutorial 2. Prerequisites: Tutorial 2 — you should be comfortable with the rollout → GAE →train_steploop,EnvPool, and thePPOConfigsurface. You do not need the off-policy tutorials (3 and 4). Time: ~20 minutes.
Every tutorial so far has quietly assumed something powerful: that the
observation the agent receives is the full state of the world. CartPole hands
you [x, x_dot, theta, theta_dot] — cart position and velocity, pole angle and
angular velocity — which is literally everything the physics needs to predict
the next step. A policy that maps the current observation to an action can be
optimal, because the current observation misses nothing. This is a fully
observable Markov Decision Process (MDP), and a memoryless (feedforward)
policy is all you ever need.
The real world is rarely so kind. Sensors drop frames, occlude, and lag; a robot sees a slice of its surroundings, not a god’s-eye state vector. When the observation is a partial, noisy shadow of the true state, you have a Partially Observable MDP (POMDP), and a memoryless policy is no longer enough — the information it needs to act well isn’t in the current frame. It has to be remembered from earlier ones. That is what this tutorial is about.
§The problem: making CartPole partially observable
CartPole-v1 is not a POMDP, so we have to make it one. The trick is to keep the
physics identical but corrupt the observation channel. FlickeringCartPole does
exactly this: on each step, with probability p (default 0.5) it blanks the
entire observed frame to all-zeros before handing it to the agent. The cart is
still there, the pole is still falling — the agent just can’t see it this step.
This is the canonical Atari-POMDP protocol from Hausknecht & Stone (2015), “Deep Recurrent Q-Learning for Partially Observable MDPs.” Blanking the whole frame (rather than adding noise) is a clean, brutal partial-observability signal: half the time, on average, the policy is acting blind.
Why does this make memory load-bearing by construction? Consider a
feedforward policy on a blanked frame. Its input is [0, 0, 0, 0] — a state
CartPole physics never produces on its own. It has no idea whether the pole is
tipping left or right, fast or slow. With no memory it can only output a fixed,
uninformed action, and it will drop the pole. A recurrent policy, by contrast,
carries a hidden state across the blank: it saw the pole tipping right two
frames ago, integrates that with its own action, and keeps balancing through the
gap. Memory is not a nice-to-have here — it is the only way to recover the
hidden state the observation withholds.
§Why velocity-masking didn’t work (a documented negative result)
The obvious first idea for a CartPole POMDP is not flickering at all — it’s to
hide the velocities, showing the agent only [x, theta] (position and
angle). That feels like it should force the agent to infer velocity from a
history of positions, which is exactly the kind of thing memory is for.
It doesn’t work, and the failure is instructive. Thrust ships MaskedCartPole
as a documented negative result: a real 500k-step run disproved it as a
POMDP. A memoryless reactive controller on [x, theta] solves it outright and
actually beats the LSTM (MLP mean return ~324 vs. LSTM ~222). The reason is
that a proportional angle/position feedback loop — “pole leaning right, push
right” — balances CartPole without ever estimating velocity. Masking the
velocities removed information the optimal policy didn’t need, so memory bought
nothing and just added variance.
The lesson: partial observability is only real if the withheld information is actually required to act well. Flickering closes the loophole that velocity-masking left open. On a blanked frame there is no reactive feedback to fall back on — position and angle are gone too — so the only way through the gap is to remember. That’s the difference between a task that merely looks like a POMDP and one that is a POMDP by construction.
§When an LSTM earns its cost
Memory isn’t free. An LSTM policy costs you, concretely:
- More parameters and compute. The recurrent trunk runs a gated update every step; training does backpropagation-through-time (BPTT) over whole trajectories, not independent transitions.
- Bookkeeping. You have to thread the hidden state
(h, c)across rollout steps, zero it at episode boundaries, and keep the training-time forward pass consistent with the behavior-time one (more on this below). Get any of that wrong and the PPO ratio silently corrupts. - Touchier optimization. Reusing a recurrent rollout across many epochs amplifies off-policy drift through the hidden state, so recurrent PPO wants fewer epochs and more careful gradient clipping than feedforward PPO.
So when does that cost pay off? Exactly when the task is a genuine POMDP —
when the information needed to act well is spread across time and missing from
any single frame. If a memoryless baseline already solves your task (as on
MaskedCartPole, or plain CartPole), an LSTM is pure overhead and often worse.
The honest way to know is to run both arms on the identical stream and compare —
which is exactly what the code below does.
§How the LSTM handles a rollout
In the feedforward loop from Tutorial 2, each step was independent: observation in, action out, nothing carried between steps. The recurrent loop threads a hidden state through:
- At each step you call
policy.forward_step(obs, state), which returns the action logits, a value estimate, and a new hidden state to feed into the next step. PassingNonefor the state starts from zeros. - On episode end, you zero the hidden state for that env. A fresh episode shares no history with the one that just ended, so its memory must start blank — otherwise stale state leaks across the boundary and poisons the new episode.
- At the boundary between rollout windows, this simple version resets the
collection state to zeros rather than carrying it across. This is a
deliberate BPTT truncation: the training-time forward pass (
evaluate_sequences) always reconstructs the hidden state from zeros at the start of each window, so the behavior policy that collected the data must do the same. If it didn’t, theold_log_probsstored during collection (computed under a warm state) would be inconsistent with thelog_probsrecomputed during training (from a cold state), and the PPO importance ratio would be garbage. Keeping both cold is the cheap, correct choice — Thrust’s design note calls this “Strategy A.”
There is one subtlety worth naming because it looks like a bug but isn’t. The
hidden-state reset and the GAE bootstrap use different notions of “episode
over.” GAE stops bootstrapping only on terminated (the pole actually fell); the
hidden-state reset (episode_starts) fires on terminated || truncated (fell
or hit the 500-step time limit). A truncated episode isn’t a real terminal
state — the pole was still up — so its value should bootstrap, but its memory
must still reset because the next episode is unrelated. The buffer and policy
handle this asymmetry internally; you just need to know it’s intentional.
§Recurrent training vs. feedforward training
Four things change from the Tutorial 2 PPO loop. If you internalize these, the code reads as “the PPO loop you know, with memory bolted on”:
- The buffer is a
RecurrentRolloutBuffer. It stores transitions trajectory-major ([n_env][T]) instead of flat, because the training forward needs whole temporally-intact sequences, not a shuffled bag of timesteps. It has an extraadd_recurrent_state(step, env, h, c)alongside the usualadd(...). - The trainer is a
RecurrentPPOTrainer, and it takes adevice. The feedforward trainer never materializes tensors on its own; the recurrent one has to build rank-3[n_env, T, obs_dim]sequence batches internally, so it needs a device handle. - You score sequences, not timesteps. The training closure calls
policy.evaluate_sequences(obs_seq, actions, initial_state, episode_starts)over rank-3 batches, replacing the rank-2evaluate_actionsfrom Tutorial 2. - Minibatches are whole trajectories. You pass
envs_per_minibatch(how many env-trajectories per minibatch) instead ofbatch_size(loose timesteps). You cannot shuffle individual timesteps without destroying the sequences the LSTM has to replay.
One more load-bearing detail that isn’t about recurrence at all:
observation normalization. CartPole’s raw observations are tiny early on
(initial perturbations ~0.05). An LSTM’s tanh gates squash such small inputs to
near-zero features, starving both the policy and value heads of gradient. So the
code below standardizes each visible coordinate to unit scale with a running
mean/std (a VecNormalize-style wrapper). The catch: a blanked frame is
exactly all-zeros, and must be passed through unchanged and excluded from the
running statistics — otherwise you’d fold zeros into the mean, pollute the
stats, and turn the “no observation” signal into a recognizable non-zero
constant. That’s why the code inlines a small ObsNormalizer with a
skip-blanked-frames special case rather than using the library’s plain
RunningMeanStd.
§The code
This trains both arms — the LSTM and the memoryless MLP — on the identical
flickering stream (same seed, same normalizer), then prints the contrast. The
budget is tiny so it runs fast in CI as a doc-test; bump TOTAL_TIMESTEPS (and
run the packaged example) for a real result. A short run won’t solve the POMDP,
so the assertion at the end only checks that both loops executed.
use std::sync::atomic::{AtomicU64, Ordering};
use burn::{
backend::{Autodiff, NdArray},
nn::LstmState,
optim::AdamConfig,
tensor::{Int, Tensor, TensorData, activation},
};
use rand::{Rng, SeedableRng, rngs::StdRng};
use thrust_rl::{
buffer::rollout::RecurrentRolloutBuffer,
env::{Environment, SpaceType, flickering_cartpole::FlickeringCartPole, pool::EnvPool},
policy::{
lstm::{LstmBurnConfig, LstmBurnPolicy},
mlp::{BurnActivation, MlpBurnConfig, MlpBurnPolicy},
},
train::{
optimizer::BurnOptimizer,
ppo::{PPOConfig, PPOTrainerBurn, RecurrentPPOTrainer},
},
};
type Backend = Autodiff<NdArray<f32>>;
// --- Hyperparameters (CI budget) -------------------------------------------
const NUM_ENVS: usize = 4; // parallel FlickeringCartPole copies
const NUM_STEPS: usize = 32; // rollout window length per env
const TOTAL_TIMESTEPS: usize = 4_096; // tiny for CI; use ~500_000 for the real run
const HIDDEN_DIM: usize = 64;
const LSTM_LR: f64 = 1.5e-3; // the POMDP wants a higher LR than plain CartPole
const MLP_LR: f64 = 3e-4; // standard CartPole PPO learning rate
const GAMMA: f32 = 0.99;
const GAE_LAMBDA: f32 = 0.95;
const N_EPOCHS: usize = 4; // low: recurrent rollouts drift off-policy fast
const ENVS_PER_MINIBATCH: usize = 2; // whole trajectories per minibatch
const FLICKER_PROB: f64 = 0.5; // blank half the frames, on average
const SEED: u64 = 0;
// CartPole's +1/step reward gives discounted returns of magnitude ~40-60.
// Against targets that large the clipped value loss (clip_range_vf) can only
// move ~0.2 per iteration and the critic freezes; scaling rewards to O(1)
// returns fixes it at the source. Reported returns stay in RAW units.
const REWARD_SCALE: f32 = 0.02;
// Train both arms on the IDENTICAL flickering stream, then print the contrast.
let (lstm_final, lstm_best) = run_lstm();
let (mlp_final, mlp_best) = run_mlp();
println!("FlickeringCartPole (p = {FLICKER_PROB}) — mean return of last <=100 episodes:");
println!(" LSTM (recurrent) : final {lstm_final:.1} best {lstm_best:.1}");
println!(" MLP (feedforward) : final {mlp_final:.1} best {mlp_best:.1}");
// With a real budget the LSTM clears the ~195 "solved" bar and the MLP plateaus
// well below it. The CI budget is far too small to solve the task, so we only
// assert both loops ran.
assert!(lstm_best >= 0.0 && mlp_best >= 0.0);
/// Build an `EnvPool` with per-env seeded flicker streams so the schedule is
/// reproducible yet decorrelated across the pool. Both arms call this with the
/// same seed, so they see the *identical* flicker stream — the only difference
/// between the runs is the policy class (memory vs. no memory).
fn make_pool() -> EnvPool<FlickeringCartPole> {
let ctr = AtomicU64::new(SEED.wrapping_mul(1000));
EnvPool::new(
|| {
let s = ctr.fetch_add(1, Ordering::Relaxed);
FlickeringCartPole::with_seed_and_probability(s, FLICKER_PROB)
},
NUM_ENVS,
)
}
/// Probe an env for its observation and (discrete) action dimensions.
fn env_dims() -> (usize, usize) {
let probe = FlickeringCartPole::new();
let obs_dim = probe.observation_space().shape[0];
let action_dim = match probe.action_space().space_type {
SpaceType::Discrete(n) => n,
_ => panic!("FlickeringCartPole is discrete"),
};
(obs_dim, action_dim)
}
/// Train the recurrent LSTM policy. Returns `(final_mean, best_mean)` return
/// over the last <=100 completed episodes (final = end of run, best = peak of
/// the moving average).
fn run_lstm() -> (f32, f32) {
let device = Default::default();
let (obs_dim, action_dim) = env_dims();
let mut env_pool = make_pool();
// Seeded LSTM policy: obs -> (action logits, value), carrying a hidden state.
let policy_config =
LstmBurnConfig { hidden_dim: HIDDEN_DIM, ..Default::default() }.with_seed(SEED);
let policy =
LstmBurnPolicy::<Backend>::with_config(obs_dim, action_dim, policy_config, &device);
let inner_opt = AdamConfig::new().init();
let burn_opt: BurnOptimizer<Backend, LstmBurnPolicy<Backend>, _> =
BurnOptimizer::new(inner_opt, LSTM_LR);
let ppo_config = PPOConfig::new()
.learning_rate(LSTM_LR)
.n_epochs(N_EPOCHS)
.gamma(GAMMA as f64)
.gae_lambda(GAE_LAMBDA as f64)
.clip_range(0.2)
// Load-bearing WITH reward scaling: a meaningful value clip keeps the
// critic tracking its (now O(1)) targets instead of freezing.
.clip_range_vf(0.2)
.vf_coef(0.5)
.ent_coef(0.01)
// Load-bearing for the recurrent trunk: without it, value-loss spikes
// wipe out the shared LSTM's policy features.
.max_grad_norm(0.5)
.target_kl(1.0); // disable KL early-stop for a short run
// The recurrent trainer takes a `device` (it materializes rank-3 batches).
let mut trainer = RecurrentPPOTrainer::new(ppo_config, policy, burn_opt, device)
.expect("valid PPO config");
let num_updates = TOTAL_TIMESTEPS / (NUM_STEPS * NUM_ENVS);
let mut norm = ObsNormalizer::new(obs_dim);
let mut buffer = RecurrentRolloutBuffer::new(NUM_STEPS, NUM_ENVS, obs_dim, HIDDEN_DIM);
let mut observations: Vec<Vec<f32>> =
env_pool.reset().iter().map(|o| norm.normalize(o)).collect();
let mut rng = StdRng::seed_from_u64(SEED);
let mut episode_returns = [0.0_f32; NUM_ENVS];
let mut completed: Vec<f32> = Vec::new();
// Collection-time recurrent state (None => zeros). Reset to None at each
// window boundary to stay consistent with the training forward (Strategy A).
let mut lstm_state: Option<LstmState<Backend, 2>> = None;
// Host copy of the state ENTERING the current step, recorded into the buffer.
let mut entering_h = vec![0.0_f32; NUM_ENVS * HIDDEN_DIM];
let mut entering_c = vec![0.0_f32; NUM_ENVS * HIDDEN_DIM];
// Host copy of the state exiting the last collected step (for warm-start).
let mut last_h = vec![0.0_f32; NUM_ENVS * HIDDEN_DIM];
let mut last_c = vec![0.0_f32; NUM_ENVS * HIDDEN_DIM];
let mut last_mean = 0.0_f32;
let mut best_mean = 0.0_f32;
for update in 0..num_updates {
buffer.reset();
for step in 0..NUM_STEPS {
let obs_flat: Vec<f32> = observations.iter().flatten().copied().collect();
let obs_t = Tensor::<Backend, 2>::from_data(
TensorData::new(obs_flat, [NUM_ENVS, obs_dim]),
&device,
);
// Thread the hidden state through the step.
let (logits, values_t, new_state) =
trainer.policy().forward_step(obs_t, lstm_state.take());
let (actions, log_probs, values_host) = sample_actions(&logits, &values_t, &mut rng);
let results = env_pool.step(&actions);
// Pull the state EXITING this step to the host so we can zero the
// rows of envs whose episode just ended.
let mut h_host: Vec<f32> = new_state.hidden.into_data().to_vec().unwrap();
let mut c_host: Vec<f32> = new_state.cell.into_data().to_vec().unwrap();
for env in 0..NUM_ENVS {
let r = &results[env];
buffer.add(
step,
env,
&observations[env],
actions[env],
r.reward * REWARD_SCALE, // train on scaled reward; report raw
values_host[env],
log_probs[env],
r.terminated,
r.truncated,
);
// Record the state that ENTERED this step (Strategy-A hook; the
// training forward recomputes from zeros, but the buffer stores
// it for completeness and future strategies).
buffer.add_recurrent_state(
step,
env,
&entering_h[env * HIDDEN_DIM..(env + 1) * HIDDEN_DIM],
&entering_c[env * HIDDEN_DIM..(env + 1) * HIDDEN_DIM],
);
episode_returns[env] += r.reward;
observations[env] = norm.normalize(&r.observation);
if r.terminated || r.truncated {
completed.push(episode_returns[env]);
episode_returns[env] = 0.0;
trainer.increment_episodes(1);
// Fresh episode => zeroed memory for this env's row.
for k in 0..HIDDEN_DIM {
h_host[env * HIDDEN_DIM + k] = 0.0;
c_host[env * HIDDEN_DIM + k] = 0.0;
}
observations[env] = norm.normalize(&env_pool.reset_env(env).unwrap());
}
}
// The state exiting this step (post per-env reset) enters the next.
entering_h.copy_from_slice(&h_host);
entering_c.copy_from_slice(&c_host);
last_h.copy_from_slice(&h_host);
last_c.copy_from_slice(&c_host);
let hidden_t = Tensor::<Backend, 2>::from_data(
TensorData::new(h_host, [NUM_ENVS, HIDDEN_DIM]),
&device,
);
let cell_t = Tensor::<Backend, 2>::from_data(
TensorData::new(c_host, [NUM_ENVS, HIDDEN_DIM]),
&device,
);
lstm_state = Some(LstmState::new(cell_t, hidden_t));
}
// Bootstrap value for the final observations under the carried state.
let last_obs_flat: Vec<f32> = observations.iter().flatten().copied().collect();
let last_obs_t = Tensor::<Backend, 2>::from_data(
TensorData::new(last_obs_flat, [NUM_ENVS, obs_dim]),
&device,
);
let boot_hidden = Tensor::<Backend, 2>::from_data(
TensorData::new(last_h.clone(), [NUM_ENVS, HIDDEN_DIM]),
&device,
);
let boot_cell = Tensor::<Backend, 2>::from_data(
TensorData::new(last_c.clone(), [NUM_ENVS, HIDDEN_DIM]),
&device,
);
let (_, last_values_t, _) = trainer
.policy()
.forward_step(last_obs_t, Some(LstmState::new(boot_cell, boot_hidden)));
let last_values: Vec<f32> = last_values_t.into_data().to_vec().unwrap();
buffer.compute_advantages(&last_values, GAMMA, GAE_LAMBDA);
// Linear LR annealing stabilizes the late-run policy above the bar.
let frac = 1.0 - (update as f64) / (num_updates.max(1) as f64);
trainer.set_learning_rate(LSTM_LR * frac.max(0.05));
// Score whole sequences, not timesteps. `None` = start each window's
// forward from a zeroed state (Strategy A); `starts` are the per-step
// episode-boundary flags the LSTM uses to reset its state mid-sequence.
let _stats = trainer
.train_step(&buffer, ENVS_PER_MINIBATCH, |p, obs_seq, actions, starts| {
p.evaluate_sequences(obs_seq, actions, None, starts)
})
.expect("recurrent train step");
// Record the cross-window carry flag, then truncate BPTT: next window's
// collection restarts from a zeroed state to match the training forward.
let final_hidden: Vec<Vec<f32>> = (0..NUM_ENVS)
.map(|e| last_h[e * HIDDEN_DIM..(e + 1) * HIDDEN_DIM].to_vec())
.collect();
let final_cell: Vec<Vec<f32>> = (0..NUM_ENVS)
.map(|e| last_c[e * HIDDEN_DIM..(e + 1) * HIDDEN_DIM].to_vec())
.collect();
buffer.seed_warm_start(NUM_STEPS - 1, &final_hidden, &final_cell);
lstm_state = None;
entering_h.iter_mut().for_each(|x| *x = 0.0);
entering_c.iter_mut().for_each(|x| *x = 0.0);
if !completed.is_empty() {
let recent = &completed[completed.len().saturating_sub(100)..];
last_mean = recent.iter().sum::<f32>() / recent.len() as f32;
best_mean = best_mean.max(last_mean);
}
}
(last_mean, best_mean)
}
/// Train the feedforward MLP baseline on the SAME flickering stream. It cannot
/// act on a blanked frame and has no memory to bridge the gaps, so it plateaus
/// below the LSTM. Returns `(final_mean, best_mean)`.
fn run_mlp() -> (f32, f32) {
let device = Default::default();
let (obs_dim, action_dim) = env_dims();
let mut env_pool = make_pool();
let policy_config = MlpBurnConfig {
num_layers: 2,
hidden_dim: 128,
use_orthogonal_init: true,
activation: BurnActivation::ReLU,
seed: Some(SEED),
};
let policy = MlpBurnPolicy::<Backend>::with_config(obs_dim, action_dim, policy_config, &device);
let inner_opt = AdamConfig::new().init();
let burn_opt: BurnOptimizer<Backend, MlpBurnPolicy<Backend>, _> =
BurnOptimizer::new(inner_opt, MLP_LR);
let ppo_config = PPOConfig::new()
.learning_rate(MLP_LR)
.n_epochs(N_EPOCHS)
.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)
.target_kl(1.0);
let mut trainer = PPOTrainerBurn::new(ppo_config, policy, burn_opt).expect("valid PPO config");
let cap = NUM_STEPS * NUM_ENVS;
let num_updates = TOTAL_TIMESTEPS / cap;
// Same normalizer as the LSTM arm so the ONLY difference is the policy class.
let mut norm = ObsNormalizer::new(obs_dim);
let mut observations: Vec<Vec<f32>> =
env_pool.reset().iter().map(|o| norm.normalize(o)).collect();
let mut episode_returns = [0.0_f32; NUM_ENVS];
let mut completed: Vec<f32> = Vec::new();
let mut last_mean = 0.0_f32;
let mut best_mean = 0.0_f32;
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);
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();
for _step in 0..NUM_STEPS {
let obs_flat: Vec<f32> = observations.iter().flatten().copied().collect();
let obs_t = Tensor::<Backend, 2>::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 in 0..NUM_ENVS {
buf_obs.extend_from_slice(&observations[env]);
buf_actions.push(actions[env]);
buf_log_probs.push(log_probs[env]);
buf_values.push(values[env]);
buf_rewards.push(results[env].reward * REWARD_SCALE);
let done = results[env].terminated || results[env].truncated;
buf_dones.push(if done { 1.0 } else { 0.0 });
episode_returns[env] += results[env].reward;
observations[env] = norm.normalize(&results[env].observation);
if done {
completed.push(episode_returns[env]);
episode_returns[env] = 0.0;
trainer.increment_episodes(1);
observations[env] = norm.normalize(&env_pool.reset_env(env).unwrap());
}
}
}
let last_obs_flat: Vec<f32> = observations.iter().flatten().copied().collect();
let last_obs_t = Tensor::<Backend, 2>::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,
);
let obs_b = Tensor::<Backend, 2>::from_data(
TensorData::new(buf_obs.clone(), [cap, obs_dim]),
&device,
);
let actions_b = Tensor::<Backend, 1, Int>::from_data(
TensorData::new(buf_actions.clone(), [cap]),
&device,
);
let old_log_probs_b =
Tensor::<Backend, 1>::from_data(TensorData::new(buf_log_probs.clone(), [cap]), &device);
let old_values_b =
Tensor::<Backend, 1>::from_data(TensorData::new(buf_values.clone(), [cap]), &device);
let advantages_b =
Tensor::<Backend, 1>::from_data(TensorData::new(advantages, [cap]), &device);
let returns_b =
Tensor::<Backend, 1>::from_data(TensorData::new(returns, [cap]), &device);
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),
)
.expect("train step");
if !completed.is_empty() {
let recent = &completed[completed.len().saturating_sub(100)..];
last_mean = recent.iter().sum::<f32>() / recent.len() as f32;
best_mean = best_mean.max(last_mean);
}
}
(last_mean, best_mean)
}
/// Host-side categorical sampling from LSTM `forward_step` outputs. `logits` is
/// `[N_env, action_dim]`, `values` is `[N_env]`. Returns
/// `(actions, action_log_probs, values)` as host `Vec`s.
fn sample_actions(
logits: &Tensor<Backend, 2>,
values: &Tensor<Backend, 1>,
rng: &mut StdRng,
) -> (Vec<i64>, Vec<f32>, Vec<f32>) {
let [n, a] = logits.dims();
let probs_host: Vec<f32> =
activation::softmax(logits.clone(), 1).into_data().to_vec().unwrap();
let log_probs_host: Vec<f32> =
activation::log_softmax(logits.clone(), 1).into_data().to_vec().unwrap();
let values_host: Vec<f32> = values.clone().into_data().to_vec().unwrap();
let mut actions = Vec::with_capacity(n);
let mut chosen_log_probs = Vec::with_capacity(n);
for i in 0..n {
let u: f32 = rng.random();
let mut cum = 0.0_f32;
let mut chosen = a - 1;
for j in 0..a {
cum += probs_host[i * a + j];
if u <= cum {
chosen = j;
break;
}
}
actions.push(chosen as i64);
chosen_log_probs.push(log_probs_host[i * a + chosen]);
}
(actions, chosen_log_probs, values_host)
}
/// Per-env GAE for the feedforward baseline. Flat `[T * N]` step-major layout
/// (index = step * num_envs + env). Identical to Tutorial 2's helper.
#[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];
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]
};
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)
}
/// Running per-dimension observation standardizer (Welford mean/variance) with
/// flicker-aware pass-through. Standardizes each *visible* coordinate to unit
/// scale (load-bearing: an LSTM's tanh gates squash CartPole's tiny raw
/// observations to near-zero features). A **blanked frame is all-zeros** — a
/// state CartPole physics never produces — so it is passed through unchanged and
/// excluded from the running statistics, keeping the "no observation" signal a
/// clean zero for both policies.
struct ObsNormalizer {
mean: Vec<f64>,
m2: Vec<f64>, // Welford's sum of squared deviations
count: f64,
}
impl ObsNormalizer {
fn new(dim: usize) -> Self {
Self { mean: vec![0.0; dim], m2: vec![0.0; dim], count: 0.0 }
}
fn normalize(&mut self, obs: &[f32]) -> Vec<f32> {
// Flickered frames are exactly all-zero; pass them through untouched.
if obs.iter().all(|&x| x == 0.0) {
return obs.to_vec();
}
self.count += 1.0;
for (i, &xf) in obs.iter().enumerate() {
let x = xf as f64;
let delta = x - self.mean[i];
self.mean[i] += delta / self.count;
self.m2[i] += delta * (x - self.mean[i]);
}
obs.iter()
.enumerate()
.map(|(i, &x)| {
let std = (self.m2[i] / self.count).sqrt().max(1e-4);
(((x as f64 - self.mean[i]) / std).clamp(-10.0, 10.0)) as f32
})
.collect()
}
}§Reading the output
Two numbers per arm, both in raw return units (undo the REWARD_SCALE in
your head — the code reports raw returns even though it trains on scaled ones):
final— mean return over the last ≤100 episodes at the end of training. Honest but noisy at episode granularity; a single unlucky late episode moves it.best— the peak of that moving average at any point during the run. This is the honest answer to “did the policy ever reach a good level within budget?”, reported alongsidefinalso end-of-run oscillation can’t hide a real success (or manufacture a fake one).
On a real run (TOTAL_TIMESTEPS ≈ 500_000), CartPole’s “solved” bar is ~195,
and the pattern you want to see is:
- LSTM
bestclears ~195 — memory recovers the hidden state through the blanks and balances the pole. - MLP
bestplateaus well below — it cannot act on a blanked frame and has no memory to bridge the gap, so it caps out far short of solved.
That gap — same env, same seed, same normalizer, only the policy class
differs — is the whole point. It’s the clean, controlled demonstration that
memory is load-bearing on this task, in a way it was not on MaskedCartPole.
The CI-sized budget in the doc-test above is far too small to reach any of this; it exists only to keep the copy-paste code compiling and running against the live API. For real numbers, run the packaged example.
§Going further
This tutorial covers i.i.d. flickering — each frame blanked independently
with probability p, the Hausknecht & Stone (2015) protocol. A few directions
from here:
-
Run the real thing. The packaged
recurrent_ppo_flickering_cartpoleexample is this exact contrast at full budget (16 envs, 128-step rollouts, 500k steps), with per-update logging and a CLI to run one arm at a time:cargo run --release --features training \ --example recurrent_ppo_flickering_cartpole -
Correlated dropout. Real sensor failures come in bursts, not independent coin flips. The
recurrent_ppo_burst_flickering_cartpoleexample blanks frames in runs of consecutive steps, which stresses memory over longer horizons. (Its burst constructor carries a stricter validity assert relatingflicker_proband burst length — out of scope here, covered in that example.) -
The design rationale.
docs/RECURRENT_POLICY_DESIGN.mdis the full design note behind the recurrent stack: why Strategy A (BPTT truncation at window boundaries) over stored-subsequence alternatives, the GAE-vs-hidden-state-reset asymmetry in detail, and the warm-start hooks the buffer exposes for future strategies.
§Next
See the tutorial index for the full dependency-ordered series.
Tutorial 6 goes the other direction — down to the metal — and shows you how to
implement the Environment trait from scratch, including the seeding and
determinism contract that EnvPool and every env in this tutorial rely on.