Skip to main content

Module recurrent_trainer

thrust_rl::train::ppo

Module recurrent_trainer 

Source
Expand description

Recurrent PPO trainer (Burn backend).

Phase 3 of the recurrent-policy epic (#262). RecurrentPPOTrainer is the rank-3 sibling of crate::train::ppo::trainer::PPOTrainerBurn: it drives the same clipped-surrogate / value-clip / entropy-bonus / KL-early-stop recipe, reusing the unchanged loss functions in crate::train::ppo::loss, but consumes whole-trajectory sequence batches from a RecurrentRolloutBuffer instead of flattened rank-2 transitions.

§Why a separate trainer

The feedforward trainer’s evaluate_fn closure is hard-typed to rank-2 observations (Tensor<B, 2>) and rank-1 actions — it flattens time and env together, structurally erasing the temporal order a recurrent policy needs. Recurrence requires a rank-3 [N_env, T, obs_dim] forward that runs the LSTM step-by-step and resets (h, c) at episode boundaries. Rather than overload the feedforward closure, this trainer takes a rank-3 analogue:

evaluate_fn(&policy, obs_seq [N_env,T,obs_dim], actions [N_env,T], episode_starts [N_env,T])
    -> (log_probs [N_env,T], entropy [N_env,T], values [N_env,T])

§Rank-2 → rank-1 shape adapter

crate::train::ppo::loss’s compute_policy_loss / compute_value_loss / compute_entropy_loss all operate on rank-1 tensors. evaluate_sequences returns rank-2 [N_env, T]. The trainer therefore flattens (reshape([N_env * T])) inside the minibatch loop, after the forward — the LSTM forward needs the [N_env, T] shape intact to run its step-by-step loop with per-step state resets. Advantage normalization runs on the flattened rank-1 view, exactly as in PPOTrainerBurn.

§Episode-boundary contract

episode_starts is consumed directly from the buffer’s materialized batch (terminated[t-1] || truncated[t-1] shifted one step, with the cross-iteration carry at t == 0). GAE stays terminated-only, computed upstream by the buffer. The two masks are intentionally distinct; this trainer does not re-derive or touch either (see docs/RECURRENT_POLICY_DESIGN.md, Q2).

§Ownership model

Identical to PPOTrainerBurn: the policy is held in Option<P> and swapped through Burn’s move-through optimizer on every gradient step.

Structs§

RecurrentPPOTrainer
Recurrent PPO trainer (Burn backend).