RL Runtime

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

Casting to Other Scalar Backends

The trust-boundary checker validates rollout JSON in host Float, because that is the interchange format. The functions below cast accepted rollouts into the element representation chosen for the proof or training path.

def TorchLean.rl.boundary.castObservation {α : Type} [Storage α] [Runtime.FromFloat α] {obsShape : Shape} (t : Tensor Float obsShape) : Tensor α obsShape

Cast a Float observation tensor into a runtime element representation α.

def TorchLean.rl.boundary.castTransition {α : Type} [Storage α] [Runtime.FromFloat α] {obsShape : Shape} {nActions : ℕ} (tr : Transition obsShape nActions) : Spec.RL.ObservedTransition (Tensor α obsShape) (Fin nActions) α

Cast a validated Float transition into a runtime element representation α.

def TorchLean.rl.boundary.castRollout {α : Type} [Storage α] [Runtime.FromFloat α] {obsShape : Shape} {nActions : ℕ} (xs : Array (Transition obsShape nActions)) : Array (Spec.RL.ObservedTransition (Tensor α obsShape) (Fin nActions) α)

Cast a whole rollout into a runtime element representation α.

def TorchLean.rl.boundary.loadRollout {obsShape : Shape} {nActions : ℕ} (path : String) (c : Contract obsShape nActions) (α : Type := Float) [Storage α] [Runtime.FromFloat α] : IO (Array (Spec.RL.ObservedTransition (Tensor α obsShape) (Fin nActions) α))

Load and validate a rollout JSON file in the requested element type.

Host Float is the default interchange representation. Select another executable or proof-facing representation with (α := ...).

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.instantiateActorCritic {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 the objective's forward execution and optimizer history.

def TorchLean.rl.ppo.trainingStep {α : 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.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.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 trainingStep keeps its optimizer history; this restores model state only.

structure TorchLean.rl.ppo.ActorCriticState (α : Type) [Storage α] (actorShapes criticShapes : List Shape) : Type

Actor and critic states, including parameters and persistent buffers.

• actor : nn.State α actorShapes

  Parameters and persistent buffers consumed by the actor graph.

• critic : nn.State α criticShapes

  Parameters and persistent buffers consumed by the critic graph.

def TorchLean.rl.ppo.splitState {σ₁ τ₁ σ₂ τ₂ : Shape} (actor : Runtime.Autograd.Model.Layers.Seq σ₁ τ₁) (critic : Runtime.Autograd.Model.Layers.Seq σ₂ τ₂) {α : Type} [Storage α] (state : nn.State α (actor.stateShapes ++ critic.stateShapes)) : ActorCriticState α actor.stateShapes critic.stateShapes

Split concatenated actor-critic state into its actor and critic components.

def TorchLean.rl.ppo.actorPolicy {obsShape logitsShape rolloutStateShape rolloutLogitsShape rolloutValueShape : Shape} {actorStateShapes : List Shape} {α : Type} [Storage α] [Context α] (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.criticValue {obsShape rolloutStateShape rolloutLogitsShape rolloutValueShape : Shape} {criticStateShapes : List Shape} {α : Type} [Storage α] [Context α] (criticGraph : nn.TypedGraphModel criticStateShapes obsShape [1] α) (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) : 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.