PPO Rollouts (Discrete Actions)

This file defines:

• fixed-horizon PPO rollout records stored as typed tensors / arrays, and
• a conversion to the minibatch format expected by the PPO autograd objective (Runtime.RL.PolicyGradient.Autograd.ppoActorCriticObjectiveDef).

The single-mask tensor GAE/return definitions live in NN.Spec.RL.Core and are re-exported by NN.Runtime.RL.Core. This typed rollout layer separates task termination from episode boundaries when computing PPO advantages.

References:

• Schulman et al., "Proximal Policy Optimization Algorithms" (2017): https://arxiv.org/abs/1707.06347
• Schulman et al., "High-Dimensional Continuous Control Using Generalized Advantage Estimation" (2015): https://arxiv.org/abs/1506.02438

Shapes

For a fixed horizon T, PPO minibatches are typically stored in "PyTorch-shaped" tensors:

• states : (T × obsShape)
• actionsOneHot : (T × nActions)
• oldLogProb : (T)
• advantages : (T)
• valueTarget : (T × 1)

@[reducible, inline]

abbrev Runtime.RL.PPO.StateBatchShape (horizon : ℕ) (obsShape : Spec.Shape) : Spec.Shape

Batch shape for a fixed-horizon sequence of observations: horizon × obsShape.

@[reducible, inline]

abbrev Runtime.RL.PPO.LogitsBatchShape (horizon nActions : ℕ) : Spec.Shape

Batch shape for a fixed-horizon sequence of action logits: horizon × nActions.

@[reducible, inline]

abbrev Runtime.RL.PPO.ScalarBatchShape (horizon : ℕ) : Spec.Shape

Batch shape for a fixed-horizon sequence of scalars: horizon.

@[reducible, inline]

abbrev Runtime.RL.PPO.ValueBatchShape (horizon : ℕ) : Spec.Shape

Batch shape for a fixed-horizon sequence of scalar values stored as a column: horizon × 1.

Rollouts

structure Runtime.RL.PPO.Step (α : Type) [TorchLean.Storage α] (obsShape : Spec.Shape) (nActions : ℕ) : Type

One fixed-horizon PPO step record.

This is the “typed parallel arrays” data layout commonly used in PPO implementations, but kept as a single record so downstream code cannot accidentally desynchronize fields.

• state : TorchLean.Tensor α obsShape

  Observation s_t (already cast into the training scalar backend).

• action : Fin nActions

  Sampled action a_t.

• oldLogProb : α

  Log-probability log π_old(a_t | s_t) under the behavior policy.

• reward : α

  Reward r_t.

• done : Bool

  Episode boundary marker (Gym-style terminated || truncated).

• value : α

  Baseline value prediction V(s_t).

• nextValue : α

  Bootstrap value prediction V(s_{t+1}) (before any auto-reset).

• terminated : Bool

  Task termination suppresses value bootstrapping. An external truncation sets done but leaves this false. The default preserves the single-mask behavior of older records.

structure Runtime.RL.PPO.Rollout (α : Type) [TorchLean.Storage α] (obsShape : Spec.Shape) (nActions horizon : ℕ) : Type

Fixed-horizon rollout buffer for PPO.

The steps_size_eq_horizon field records the invariant that the buffer has exactly horizon steps; this lets downstream tensor conversion be total without runtime bounds checks.

• steps : Array (Step α obsShape nActions)
• steps_size_eq_horizon : self.steps.size = horizon

  Invariant: fixed-horizon rollouts always have exactly horizon steps.

structure Runtime.RL.PPO.TrainingBatch (α : Type) [TorchLean.Storage α] (obsShape : Spec.Shape) (nActions horizon : ℕ) : Type

Named tensors consumed by one PPO actor-critic update.

• states : TorchLean.Tensor α (StateBatchShape horizon obsShape)

  Observations for each rollout step.

• actionsOneHot : TorchLean.Tensor α (LogitsBatchShape horizon nActions)

  Sampled actions encoded against the action-logit axis.

• oldLogProb : TorchLean.Tensor α (ScalarBatchShape horizon)

  Behavior-policy log-probability for each sampled action.

• advantages : TorchLean.Tensor α (ScalarBatchShape horizon)

  Normalized generalized advantages used by the policy objective.

• valueTargets : TorchLean.Tensor α (ValueBatchShape horizon)

  Lambda-return targets used by the value objective.

def Runtime.RL.PPO.TrainingBatch.Internal.arguments {α : Type} [TorchLean.Storage α] {obsShape : Spec.Shape} {nActions horizon : ℕ} (batch : TrainingBatch α obsShape nActions horizon) : TorchLean.TensorPack α [StateBatchShape horizon obsShape, LogitsBatchShape horizon nActions, ScalarBatchShape horizon, ScalarBatchShape horizon, ValueBatchShape horizon]

Pack a named PPO batch for the low-level autograd objective.

def Runtime.RL.PPO.Rollout.Internal.generalizedAdvantageEstimationWithBoundaries {α : Type} [TorchLean.Storage α] [Context α] {n : ℕ} (gamma lam : α) (rewards values nextValues : TorchLean.Tensor α [n]) (terminated boundaries : TorchLean.Tensor Bool [n]) : TorchLean.Tensor α [n]

GAE with separate masks for the next-state value and continuation into the next step.

def Runtime.RL.PPO.Rollout.generalizedAdvantages {α : Type} [TorchLean.Storage α] [Context α] {obsShape : Spec.Shape} {nActions horizon : ℕ} (gamma lam : α) (r : Rollout α obsShape nActions horizon) : TorchLean.Tensor α [horizon]

Unnormalized GAE for a PPO rollout. Task termination suppresses the next-state value; every episode boundary stops advantage continuation. Thus a truncation bootstraps from nextValue before auto-reset without using rewards from the following episode.

def Runtime.RL.PPO.Rollout.trainingBatch {α : Type} [TorchLean.Storage α] [Context α] {obsShape : Spec.Shape} {nActions horizon : ℕ} [NeZero horizon] [NeZero nActions] (gamma lam : α) (r : Rollout α obsShape nActions horizon) : IO (TrainingBatch α obsShape nActions horizon)

Convert a fixed-horizon rollout into the PPO minibatch expected by Autograd.ppoActorCriticObjectiveDef.

Notes:

• Advantages are normalized (z-score) for the policy-gradient term, a common PPO variance-reduction practice. Value targets (lambda-returns) are computed from the unnormalized advantages.
• Termination suppresses bootstrapping; truncation retains the pre-reset next-state value. Both stop advantage continuation across the episode boundary.