PPO on Lean-native GridWorld (Executable Demo + Formal Model)

This example complements torchlean ppo_cartpole:

• torchlean ppo_cartpole uses an external Python Gymnasium environment and checks every step against a Lean-side trust-boundary contract (Runtime.RL.Boundary.Contract).
• This example uses a Lean-native GridWorld and still runs the same PPO update in Lean.

Even though the environment is defined in Lean, we still validate every transition with the boundary checker. That keeps the data model unified: downstream training code consumes Spec.RL.ObservedTransition in a single format regardless of whether the source is a Lean-native environment or an external sampler.

Formal hooks

1. The environment has an induced finite stochastic MDP (Spec.RL.FiniteStochastic.MDP) and we import a proof that it is well-formed (row-stochastic transition rows, 0 ≤ γ < 1).
2. The boundary checker can be turned into a Prop-level hypothesis via Proofs.RL.Boundary.contractHolds_of_checkTransitionFin_eq_ok (see NN/Proofs/RL/Boundary.lean), or you can use the proof-layer Gymnasium wrapper Runtime.RL.Gymnasium.Session.stepCheckedWithProof (NN/Proofs/RL/Gymnasium.lean) for external environments.

CLI flags

• --cuda: run the Torch backend on CUDA (requires building with -K cuda=true).
• --updates <n>: number of PPO updates to run.
• --eval-every <n>: evaluate the greedy policy every n updates.
• --eval-episodes <n>: number of evaluation episodes per checkpoint.
• --eval-max-steps <n>: maximum steps per evaluation episode.
• --log <path>, --policy <path>, --path <path>: override artifact output paths.

Run (from the repo root):

lake exe torchlean ppo_gridworld
lake build -R -K cuda=true && lake exe torchlean ppo_gridworld --cuda
lake exe torchlean ppo_gridworld --updates 200

Artifacts:

• The executable writes widget-friendly JSON snapshots to data/rl/ by default: ppo_gridworld_trainlog.json, ppo_gridworld_policy.json, ppo_gridworld_path.json (override with --log, --policy, --path).
• You can also tune runtime cost with: --updates, --eval-every, --eval-episodes, --eval-max-steps.
• Visualize them in the editor via NN/Examples/RL/PPOGridWorldView.lean.

What this example does (and does not) guarantee

• Because the environment dynamics are Lean code, you can reason about its properties directly (e.g. determinism, Markov property w.r.t. the explicit state, bounded rewards).
• The PPO/GAE update is implemented as Lean definitions and a TorchLean autograd program, so it is a natural target for formal proofs about the update equation.
• As in most practical PPO code, convergence and optimality are not guaranteed by this example; it is tuned for inspectability and type safety, not leaderboard performance.

References (primary):

• 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
• Williams, "Simple statistical gradient-following algorithms for connectionist reinforcement learning" (REINFORCE, 1992): https://doi.org/10.1007/BF00992696
• Sutton and Barto, Reinforcement Learning: An Introduction (2nd ed., GridWorld examples): http://incompleteideas.net/book/the-book-2nd.html
• Puterman, Markov Decision Processes (finite discounted MDPs): https://doi.org/10.1002/9780470316887

def NN.Examples.Models.RL.PPOGridWorld.exeName : String

Name of this executable target (used in CLI error messages and banners).

Configuration

def NN.Examples.Models.RL.PPOGridWorld.width : ℕ

Grid width (number of columns).

def NN.Examples.Models.RL.PPOGridWorld.height : ℕ

Grid height (number of rows).

def NN.Examples.Models.RL.PPOGridWorld.nStates : ℕ

Total number of discrete states (width * height).

def NN.Examples.Models.RL.PPOGridWorld.nActions : ℕ

Number of discrete actions (up/down/left/right).

def NN.Examples.Models.RL.PPOGridWorld.hiddenDim : ℕ

Width of the hidden layer in the actor and critic MLPs.

def NN.Examples.Models.RL.PPOGridWorld.horizon : ℕ

PPO rollout horizon (also the training batch size for this example).

def NN.Examples.Models.RL.PPOGridWorld.gamma : Float

Discount factor used in returns / GAE.

def NN.Examples.Models.RL.PPOGridWorld.lam : Float

GAE(λ) parameter controlling the bias/variance tradeoff of advantage estimates.

def NN.Examples.Models.RL.PPOGridWorld.lr : Float

Adam learning rate.

def NN.Examples.Models.RL.PPOGridWorld.updateEpochs : ℕ

Number of PPO optimization epochs per collected rollout batch.

def NN.Examples.Models.RL.PPOGridWorld.updatesMax : ℕ

Default maximum number of PPO updates (can be overridden by --updates).

def NN.Examples.Models.RL.PPOGridWorld.evalEvery : ℕ

Default evaluation checkpoint interval (can be overridden by --eval-every).

def NN.Examples.Models.RL.PPOGridWorld.evalEpisodes : ℕ

Default evaluation episodes per checkpoint (can be overridden by --eval-episodes).

instance NN.Examples.Models.RL.PPOGridWorld.instFactLtNatOfNatHorizon : Fact (0 < horizon)

instance NN.Examples.Models.RL.PPOGridWorld.instFactLtNatOfNatNActions : Fact (0 < nActions)

def NN.Examples.Models.RL.PPOGridWorld.obsShape : Shape

The observation tensor shape used by this example: [..., nStates] one-hot vectors.

def NN.Examples.Models.RL.PPOGridWorld.pfxBatch : Shape

def NN.Examples.Models.RL.PPOGridWorld.sStateBatch : Shape

def NN.Examples.Models.RL.PPOGridWorld.sLogitsBatch : Shape

def NN.Examples.Models.RL.PPOGridWorld.sScalarBatch : Shape

def NN.Examples.Models.RL.PPOGridWorld.sValueBatch : Shape

def NN.Examples.Models.RL.PPOGridWorld.sState1 : Shape

def NN.Examples.Models.RL.PPOGridWorld.sLogits1 : Shape

def NN.Examples.Models.RL.PPOGridWorld.sValue1 : Shape

Formal GridWorld model (spec/proof layer)

We define a real-valued GridWorld model and record the proof that its stochastic-MDP view is valid.

This proof is not used by the executable training loop directly; it exists so that downstream theorems about the induced MDP can refer to a concrete environment used in an example.

def NN.Examples.Models.RL.PPOGridWorld.startPos : Spec.RL.Envs.GridWorld.State width height

Start position (top-left cell).

def NN.Examples.Models.RL.PPOGridWorld.goalPos : Spec.RL.Envs.GridWorld.State width height

Goal position (bottom-right cell).

noncomputable def NN.Examples.Models.RL.PPOGridWorld.discountR : ℝ

A discount factor in [0,1) at the proof layer (ℝ), used to build an MDP instance.

noncomputable def NN.Examples.Models.RL.PPOGridWorld.gwR : Spec.RL.Envs.GridWorld width height

Proof-layer GridWorld instance over ℝ rewards/discounts.

theorem NN.Examples.Models.RL.PPOGridWorld.gwR_valid : Spec.RL.FiniteStochastic.Valid gwR.toFiniteStochasticMDP

The induced finite stochastic MDP for gwR is well-formed (0 ≤ γ < 1, row-stochastic transitions).

Lean-native runtime environment

We implement a Gym-style environment (Spec.RL.Env) whose observations are one-hot vectors over the flattened finite state space Fin nStates.

This keeps the PPO code identical to the Gymnasium example: the policy consumes tensors and produces logits over a finite action set.

def NN.Examples.Models.RL.PPOGridWorld.startState : Fin nStates

Start state encoded as Fin nStates.

def NN.Examples.Models.RL.PPOGridWorld.goalState : Fin nStates

Goal state encoded as Fin nStates.

def NN.Examples.Models.RL.PPOGridWorld.obsOfState (s : Fin nStates) : Spec.Tensor Float obsShape

Observation function: encode the discrete state as a one-hot vector of length nStates.

def NN.Examples.Models.RL.PPOGridWorld.coordDist (a b : ℕ) : Float

Absolute difference on natural-number coordinates, returned as a Float.

def NN.Examples.Models.RL.PPOGridWorld.goalDistance (pos : Spec.RL.Envs.GridWorld.State width height) : Float

Manhattan distance to the goal.

def NN.Examples.Models.RL.PPOGridWorld.stepState (state : Fin nStates) (action : Fin nActions) : Spec.RL.StepResult (Fin nStates) Float

Deterministic GridWorld transition function with dense progress rewards.

The original sparse -1 until terminal reward gave short runs too little learning signal: random rollouts rarely found the goal, so PPO received almost no useful signal. This shaped reward keeps the same goal-reaching task, but gives the learner immediate credit for moving closer to the goal and a small penalty for dithering.

def NN.Examples.Models.RL.PPOGridWorld.env : Spec.RL.Env (Fin nStates) (Fin nActions) (Spec.Tensor Float obsShape) Float

Lean-native environment packaged as a Spec.RL.Env for reuse with the generic RL runtime.

Trust boundary contract

Even though this environment is Lean-native, we keep a contract in play to exercise the “checked preconditions” workflow and to keep the interface identical to the external Gymnasium collector.

def NN.Examples.Models.RL.PPOGridWorld.contract : Runtime.RL.Boundary.Contract obsShape nActions

Model (Actor + Critic)

def NN.Examples.Models.RL.PPOGridWorld.modelCfg : API.nn.models.PPOActorCriticConfig

def NN.Examples.Models.RL.PPOGridWorld.actorMk (pfx : Shape) : API.nn.M (API.nn.Sequential (Spec.Shape.appendDim pfx nStates) (Spec.Shape.appendDim pfx nActions))

Construct the actor network as an MLP mapping one-hot observations to action logits.

def NN.Examples.Models.RL.PPOGridWorld.criticMk (pfx : Shape) : API.nn.M (API.nn.Sequential (Spec.Shape.appendDim pfx nStates) (Spec.Shape.appendDim pfx 1))

Construct the critic network as an MLP mapping one-hot observations to a scalar value estimate.

Rollout collection (Lean-native environment)

def NN.Examples.Models.RL.PPOGridWorld.collectRolloutNativeWith (predictLogits : Spec.Tensor Float obsShape → Spec.Tensor Float (Tensor.shapeOfDims [nActions])) (predictValue : Spec.Tensor Float obsShape → Float) (rngSeed rngCounter : ℕ) (resetOnDone : Bool := true) : IO (Runtime.RL.PPO.Rollout Float obsShape nActions horizon × ℕ)

Collect a fixed-horizon PPO rollout from the Lean-native environment using a checked session.

This is an example-local wrapper around rl.ppo.collectRolloutCheckedSessionWith that packages:

• a Spec.RL.Env as a rl.session.CheckedSession, and
• the (actor, critic) prediction functions at observation shape.

Evaluation

Evaluation helpers live in NN.API.rl.eval (runtime module NN.Runtime.RL.Eval).

Main Training Loop

def NN.Examples.Models.RL.PPOGridWorld.main (args : List String) : IO UInt32

Entry point for lake exe torchlean ppo_gridworld.

This executable:

• runs PPO updates against a Lean-native GridWorld environment,
• periodically evaluates the greedy policy and logs the average return,
• writes widget-friendly JSON artifacts (training curve, greedy policy snapshot, greedy path snapshot).