RL Runtime

Rollout boundary checks, Gymnasium sessions, Float32 and interval numerics, and PPO actor-critic wiring exposed under TorchLean.rl.

structure TorchLean.rl.ppo.ActorCritic (α : Type) [Storage α] [Context α] (stateShapes : List Shape) (stateShape : Shape) (batch nActions : ℕ) : Type

PPO runtime state. Stateful layers update their buffers from the actual training forward pass.

• objective : TorchLean.Module.Objective α Unit stateShapes [stateShape, [batch, nActions], [batch], [batch], [batch, 1]]

def TorchLean.rl.ppo.ActorCritic.create {stateShape : Shape} {batch nActions : ℕ} {α : Type} [Storage α] [Context α] [Runtime.FromFloat α] [Runtime.Autograd.Torch.TensorTransfer α] (options : Runtime.Autograd.Torch.Config) (actor : Runtime.Autograd.Model.Layers.Seq stateShape [batch, nActions]) (critic : Runtime.Autograd.Model.Layers.Seq stateShape [batch, 1]) : IO (ActorCritic α (actor.stateShapes ++ critic.stateShapes) stateShape batch nActions)

Instantiate the standard PPO actor-critic runtime.

The actor and critic share one objective. Use ActorCritic.bind to attach an optimizer and obtain an update callback with persistent optimizer history.

def TorchLean.rl.ppo.ActorCritic.bind {α : Type} [Storage α] [Context α] [Runtime.FromFloat α] {stateShapes : List Shape} {obsShape : Shape} {batch nActions : ℕ} (m : ActorCritic α stateShapes (StateBatchShape batch obsShape) batch nActions) (config : optim.Optimizer) : IO (TrainingBatch α obsShape nActions batch → IO Unit)

Bind a PPO actor-critic update function and preserve its optimizer history across calls.

Each call updates model buffers from the activations used to compute the gradients, then performs the optimizer step. This includes every repeated PPO epoch over the same rollout batch.

def TorchLean.rl.ppo.ActorCritic.state {α : Type} [Storage α] [Context α] {stateShapes : List Shape} {stateShape : Shape} {batch nActions : ℕ} (m : ActorCritic α stateShapes stateShape batch nActions) : IO (nn.State α stateShapes)

Read concatenated actor-critic state without refreshing buffers.

def TorchLean.rl.ppo.ActorCritic.setState {α : Type} [Storage α] [Context α] {stateShapes : List Shape} {stateShape : Shape} {batch nActions : ℕ} (m : ActorCritic α stateShapes stateShape batch nActions) (newState : nn.State α stateShapes) : IO Unit

Restore actor and critic parameters and persistent buffers.

An already-bound ActorCritic.bind callback keeps its optimizer history; this restores model state only.

def TorchLean.rl.ppo.policy {obsShape logitsShape rolloutStateShape rolloutLogitsShape rolloutValueShape : Shape} {actorStateShapes : List Shape} {α : Type} [Storage α] (actorGraph : nn.TypedGraphModel actorStateShapes obsShape logitsShape α) (actorRollout : nn.Sequential rolloutStateShape rolloutLogitsShape) (criticRollout : nn.Sequential rolloutStateShape rolloutValueShape) (state : nn.State α (Runtime.Autograd.Model.Layers.Seq.stateShapes actorRollout ++ Runtime.Autograd.Model.Layers.Seq.stateShapes criticRollout)) (sameActorState : Runtime.Autograd.Model.Layers.Seq.stateShapes actorRollout = actorStateShapes := by rfl) : Tensor α obsShape → Tensor α logitsShape

Build a single-observation actor policy from the state of a rollout-shaped actor-critic module.

The typed actor graph records its state layout, while sameActorState states that the rollout actor uses that layout as well.

def TorchLean.rl.ppo.value {obsShape valueShape rolloutStateShape rolloutLogitsShape rolloutValueShape : Shape} {criticStateShapes : List Shape} {α : Type} [Storage α] (criticGraph : nn.TypedGraphModel criticStateShapes obsShape valueShape α) (actorRollout : nn.Sequential rolloutStateShape rolloutLogitsShape) (criticRollout : nn.Sequential rolloutStateShape rolloutValueShape) (state : nn.State α (Runtime.Autograd.Model.Layers.Seq.stateShapes actorRollout ++ Runtime.Autograd.Model.Layers.Seq.stateShapes criticRollout)) (sameCriticState : Runtime.Autograd.Model.Layers.Seq.stateShapes criticRollout = criticStateShapes := by rfl) (oneValue : valueShape.size = 1 := by decide) : Tensor α obsShape → α

Build a single-observation critic function from the state of a rollout-shaped actor-critic module.

The result is scalar because the typed critic graph has a checked one-element output shape. Scalar tensors and any number of singleton axes are accepted.