PPO Rollout Collection (Checked Sessions)

This file provides the rollout-collection loop used by executable PPO workflows. The key goals are:

• keep data collection typed and total (no “parallel arrays” that can desync),
• enforce the trust-boundary contract on every step (external Gymnasium or Lean-native env), and
• keep the API usable: callers should not need to thread a dozen actor/critic lowering details through every function call.

The unified session interface lives in NN.Runtime.RL.Session (Session.CheckedSession). The lower-level Gymnasium subprocess protocol is implemented in NN.Runtime.RL.Gymnasium.

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
• Gymnasium API reference (reset/step, terminated vs truncated): https://gymnasium.farama.org/

Rollout collection (ergonomic core API)

def Runtime.RL.PPO.collect {α : Type} [TorchLean.Storage α] [Context α] {obsShape : Spec.Shape} {nActions horizon : ℕ} [NeZero horizon] [NeZero nActions] (session : Session.CheckedSession obsShape nActions) (castObservation castReward : Float → α) (predictLogits : TorchLean.Tensor α obsShape → TorchLean.Tensor α [nActions]) (predictValue : TorchLean.Tensor α obsShape → α) (rngSeed rngCounter : ℕ) : IO (Rollout α obsShape nActions horizon × ℕ)

Collect a fixed-horizon rollout from any stateful environment session that can produce fully-observed, contract-checked transitions.

The caller provides:

• session: initialization, observation, and checked stepping for the environment,
• castObservation to inject host Float observations into the chosen scalar backend α,
• castReward to inject host Float rewards into the chosen scalar backend α,
• predictLogits for the current actor,
• predictValue for the current critic (returns a scalar α).

The API supports the “typed graph + parameters” calling convention used throughout TorchLean. Stored log-probabilities use logSoftmax without clamping. When using this rollout with Autograd.PPO.objective, selected log-probabilities must lie inside its [-10^30, 10^30] clamp and the policy reductions must agree for identical parameters to yield a unit ratio.

Rollout collection from Gymnasium (subprocess bridge)

def Runtime.RL.PPO.collectRolloutFromGymnasium {α : Type} [TorchLean.Storage α] [Context α] {obsShape : Spec.Shape} {nActions horizon : ℕ} [NeZero horizon] [NeZero nActions] (castObservation castReward : Float → α) (gym : Gymnasium.Client obsShape nActions) (predictLogits : TorchLean.Tensor α obsShape → TorchLean.Tensor α [nActions]) (predictValue : TorchLean.Tensor α obsShape → α) (rngSeed rngCounter resetSeed : ℕ) : IO (Rollout α obsShape nActions horizon × ℕ)

Collect a fixed-horizon rollout from a Gymnasium subprocess environment.

This specializes collect to Gymnasium.Session.