PPO GridWorld Run🔗

A small CPU run connects the environment, rollout, and update. The environment is a 4\times4 GridWorld defined in Lean. The actor and critic use a fixed rollout horizon of 64. Run one update and write the artifacts to files:

# Keep the evaluation path and policy so the measured return
# can be reconstructed.
lake exe torchlean ppo_gridworld --device cpu \
  --updates 1 \
  --eval-every 1 --eval-episodes 1 --eval-max-steps 8 \
  --log /tmp/ppo-gridworld-trainlog.json \
  --policy /tmp/ppo-gridworld-policy.json \
  --path /tmp/ppo-gridworld-path.json

A captured run produced the following output and exited 0:

[TorchLean] arithmetic: native binary32
[TorchLean] execution: eager
[TorchLean] device: cpu
ppo_gridworld: PPO on Lean-native GridWorld (4x4, horizon=64) (device=cpu)
  env: pure Lean dynamics + boundary contract check + formal MDP validity proof available
  eval(step=0) avg_return=-0.400000
  update=1 avg_return=3.600000
  wrote TrainLog JSON: /tmp/ppo-gridworld-trainlog.json
ppo_gridworld: wrote policy snapshot to /tmp/ppo-gridworld-policy.json
ppo_gridworld: wrote path snapshot to /tmp/ppo-gridworld-path.json
ppo_gridworld: done
ppo_gridworld: ok

The transcript records four kinds of information:

1. the three [TorchLean] banners are configuration facts, described in the command-line reference;

2. update=1 is an execution fact: the PPO program completed one rollout/update cycle without tripping a contract check; the configured optimization epochs reuse that collected batch;

3. the two avg_return numbers are measurements from one seeded run, not properties of PPO;

4. the env: line advertises that a formal MDP validity proof exists for this environment. That proof is a theorem about the specification, and we look at it directly later in this chapter.

The implementation is NN/Examples/Models/RL/PPOGridWorld.lean, and the pure environment is NN.Spec.RL.Envs.GridWorld.

GridWorld is useful here because an observation has a concrete interpretation: one of sixteen cells, encoded as a one-hot vector. Four actor logits select a direction, and the next cell can be calculated without an external simulator. That makes the policy's behavior recoverable from the saved path, rather than leaving a scalar reward curve as the only evidence of what happened. The horizon of sixty-four concerns training collection; the eight-step evaluation cap concerns a separate greedy episode. They need not have the same length.

The command uses eight optimization epochs per collected rollout. Thus --updates 1 collects one batch and reuses it for the configured optimization passes. The displayed counter is not a count of individual environment actions or a promise of one gradient evaluation. This distinction matters when comparing a longer rollout with more update epochs: one changes the collected data, while the other changes how often the existing data influences the parameters.

Episode Returns🔗

The evaluation metric is rl.eval.averageEpisodeTotalReward, the undiscounted sum of rewards over a greedy episode, capped at --eval-max-steps. The command also writes the greedy policy and path. Together with the reward definition, these artifacts let us recompute the two reported returns.

The path artifact records decoded (row, column) states:

before: [[0,0],[1,0],[1,1],[0,1],[0,0],[1,0],[1,1],[0,1],[0,0]]
after:  [[0,0],[0,1],[1,1],[1,2],[1,3],[1,3],[1,3],[1,3],[1,3]]

The untrained greedy policy walks a four-step cycle and comes back to the start twice. The updated one moves toward the goal for four steps and then stalls against the right wall. The policy artifact explains both, once you know that actions are encoded 0=up, 1=down, 2=left, 3=right, and that state index row * 4 + column is the row-major flattening:

before: [1, 2, 0, 0, 3, 0, 3, 3, 3, 0, 1, 3, 2, 1, 0, 3]
after:  [3, 1, 0, 3, 3, 3, 3, 3, 3, 3, 1, 3, 2, 1, 0, 3]

Read the after row along the visited cells: cell 0 says right, cell 1 says down, cell 5 says right, cell 6 says right, cell 7 says right. That is exactly the recorded path, and the final right at cell 7 is why the episode stalls: column 3 is the last column, and this GridWorld clamps at its borders instead of wrapping or failing.

The path contains nine states because eight transitions need both a starting state and a successor for every action. The policy snapshot instead gives an action for every cell, including cells the path never visits. Reading them together checks that the chosen actions explain the observed transitions. It also avoids interpreting a larger return as reaching the goal: the last four repeated coordinates show exactly what the updated greedy policy still fails to do.

-- Reward distance reduction, then charge a small cost for
-- spending another step.
let progress := goalDistance pos - goalDistance nextPos
{ state := nextState
  reward := progress - 0.05
  terminated := false
  truncated := false }

[(18 : Rat)/5, (53 : Rat)/20, (17 : Rat)/10, (3 : Rat)/4, (-1 : Rat)/5, (-3 : Rat)/20, (-1 : Rat)/10, (-1 : Rat)/20]

[(-2 : Rat)/5,
 (-27 : Rat)/20,
 (-23 : Rat)/10,
 (-5 : Rat)/4,
 (-1 : Rat)/5,
 (-23 : Rat)/20,
 (-21 : Rat)/10,
 (-21 : Rat)/20]

GridWorld Definition🔗

NN.Spec.RL.Envs.GridWorld defines the state, actions, movement, and sparse reward in Lean. The MDP theorems use these definitions directly.

The state is a pair of Fin coordinates and an action is one of four directions. Movement is border-clamped rather than wrapping or erroring, which is the behaviour we saw stall the trained policy at column 3:

(0, 0)

(1, 3)

Moving up from the top row leaves you in the top row, and moving right from the last column leaves you there. This is a deliberate choice: a clamped grid is a total function, so the environment needs no error case, and the induced MDP has no absorbing failure state to reason about.

To hand the environment to the finite MDP layer, coordinates are flattened row-major through mathlib's finProdFinEquiv (Community, 2020):

(7, 1, 3)

The mathlib equivalence supplies both the encoding and its inverse laws, so subsequent proofs can use the flattening as a bijection.

Now the reward. Here is the specification environment, instantiated at the grid the examples use, together with the reward on a step that does not reach the goal:

The reward has type ℝ, so the block proves an equality rather than evaluating a float. The specification assigns exactly -1; rounding is not part of this statement.

After unfolding the step definition, norm_num discharges (0,0)\neq(3,3) and (0,1)\neq(3,3) using Fin.ext_iff and reads off the reward field. A change to the reward convention would require revisiting this proof.

The sparse reward theorem and the shaped training reward share movement rules, but they do not define the same value function. A value is an expected accumulated reward under a particular reward convention and discount. Changing progress rewards can therefore change which policy is preferred unless an appropriate invariance result is supplied. Here the explicit distinction lets us use the GridWorld theorem for its stated MDP without silently transferring an optimality claim to the shaped PPO experiment.

Actor And Critic Batch Shapes🔗

The GridWorld command does not hide its actor and critic inside the training loop. It starts from a small reusable configuration, which nn.models.PPO turns into a two-layer tanh MLP per role: the actor ends in four action logits, the critic in one value.

The constructors take a leading batch shape, so the same configuration describes both individual observations and batches:

([16], [4], [1])

([64, 16], [64, 4], [64, 1])

Scoring a single state during a rollout and a whole horizon during optimization uses the same model at two batch shapes. PyTorch's nn.Linear(16, 4) accepts leading dimensions implicitly (Paszke et al., 2019). Here the prefix belongs to the sequential model's type, so the model layer can state and check it.

The configuration also validates, which matters because these numbers usually arrive from flags:

Except.ok ()

Except.error "PPO: action count must be positive"

A zero action count would otherwise build a model whose logit axis has length zero, and the first symptom would be a categorical sampler that cannot return anything.

GridWorld, CartPole, and Pong RAM reuse these two constructors with different observation and action widths. Environment collection and boundary checking stay outside the model helper. The actor definition can therefore serve both a single-state rollout and a full-horizon update.

The actor's four outputs are logits; their differences determine the action distribution after normalization. The critic's one output estimates future reward in the units of the chosen reward convention. These heads have different training targets even though they share the same observation width and hidden-width configuration. The [64, 1] critic output must eventually align with one scalar target per collected step; the singleton axis is a model interface choice, not sixty-four additional value predictions.

Rollout Data🔗

For a discrete action space of size A, one PPO step stores

(s_t,\;a_t,\;\log\pi_{\mathrm{old}}(a_t\mid s_t),\; r_t,\;d_t,\;\mathrm{terminated}_t,\;V(s_t),\;V(s_{t+1})).

The Lean structure in NN.Runtime.RL.PPO.Rollout holds one observation tensor, one Fin nActions action, four scalars, and two Boolean markers per step. done marks either termination or truncation and stops advantage continuation. terminated suppresses the next-state value bootstrap. When constructing a step manually, omitting terminated defaults it to done, preserving the earlier single-mask behavior. The rollout container also fixes the number of steps:

structure Rollout (α : Type) [TorchLean.Storage α]
    (obsShape : Shape) (nActions horizon : Nat) where
  steps : Array (Step α obsShape nActions)
  /-- Invariant: fixed-horizon rollouts always have exactly `horizon` steps. -/
  steps_size_eq_horizon : steps.size = horizon

The array carries its length as a proof field. With separate arrays for rewards, values, and observations, an extra entry after an early episode end can misalign the fields. A single step record keeps its fields together, and the proof fixes the number of records. Neither prevents a caller from putting semantically mismatched observations and rewards into a step.

Rollout.trainingBatch converts the records into five named tensors whose shapes are computed from the horizon:

The runtime's shape abbreviations compute these dimensions:

([64, 16], [64, 4], [64], [64, 1])

The conversion from rollout to batch is total: because the horizon invariant is already in the structure, no step of the packing can fail, and trainingBatch needs no error case at all. The one subtlety is deliberate and documented in the source: value targets are computed from the unnormalized advantages, and only the policy term sees z-score normalized advantages. Normalizing the advantages that feed the value target would change what the critic regresses to, which is a different algorithm than the one Schulman et al. (2017) describe.

The stored old log probability belongs to the policy that sampled the action. Recomputing it after an optimizer update would change the denominator of the importance ratio and erase the comparison PPO is trying to make. Likewise, nextValue refers to the successor before any automatic environment reset. A correctly sized rollout could still be wrong if either field were filled from the wrong policy version or episode. The record keeps related fields adjacent; collection determines their semantic alignment.

Returns And Generalized Advantage Estimation🔗

Let d_t mark an episode boundary and z_t mark task termination. Both are zero or one. An external time-limit truncation has d_t=1 and z_t=0. The one-step temporal difference residual used by PPO is

\delta_t =r_t+\gamma(1-z_t)V(s_{t+1})-V(s_t),

generalized advantage estimation runs the backward recursion

A_t =\delta_t+\gamma\lambda(1-d_t)A_{t+1},

and the matching value target is R_t=A_t+V(s_t) (Schulman et al., 2015).

The generic rl.core.generalizedAdvantageEstimation API below has one mask and uses it in both places. These examples have no external truncations, so z_t=d_t. PPO's Rollout.generalizedAdvantages keeps the two masks separate.

Begin with discounted returns at \gamma=1/2 over rewards 1,2,3:

[(11 : Rat)/4, (7 : Rat)/2, 3]

By hand: G_2=3, G_1=2+\tfrac12\cdot3=\tfrac72, G_0=1+\tfrac12\cdot\tfrac72=\tfrac{11}4. The exact fractions are the reason to run this at ℚ first. At Float the same call gives the decimal form, and here the two agree because every value is a dyadic rational:

[2.750000, 3.500000, 3.000000]

A done flag stops the return recursion at an episode boundary:

[2, 2, 3]

Step 1 terminates, so its return is its own reward, 2, with no \tfrac12\cdot3 attached. Step 0 still bootstraps from step 1, giving 1+\tfrac12\cdot2=2. Compare with [11/4, 7/2, 3] above: one Boolean changed two of the three targets. An incorrect boundary flag would silently change these targets while leaving every tensor shape intact.

Now generalized advantage estimation. With \lambda=1 and a zero baseline it must degenerate to the Monte Carlo return, which is a good self-check on the recursion direction:

[(11 : Rat)/4, (7 : Rat)/2, 3]

This matches discountedReturns. The next example changes both the baseline values and \lambda, setting \lambda=1/2. Its smaller advantages reflect both changes:

[(11 : Rat)/8, (3 : Rat)/2, 0]

The four tensors contain rewards, values, next values, and termination flags, in that order. Following the recursion: \delta_2=3+0-3=0 because the step is terminal, \delta_1=2+\tfrac12\cdot3-2=\tfrac32 and A_1=\tfrac32, then \delta_0=1+\tfrac12\cdot2-1=1 and A_0=1+\tfrac14\cdot\tfrac32=\tfrac{11}8. Adding the baseline back recovers the value targets:

[(19 : Rat)/8, (7 : Rat)/2, 3]

The parameter \lambda changes how far a temporal-difference residual propagates backward. With \gamma=\lambda=1/2, each additional step multiplies that contribution by 1/4. This is distinct from discounting the raw reward by 1/2: the residual has already compared a reward-plus-bootstrap prediction against the current baseline. In the final terminal step here, the baseline equals the reward, so its residual and advantage are zero even though the reward itself is three. Adding the baseline back restores a value target of three.

# Carry only the following advantage backwards through the
# episode mask.
rewards     = torch.tensor([1.0, 2.0, 3.0])
values      = torch.tensor([1.0, 2.0, 3.0])
next_values = torch.tensor([2.0, 3.0, 0.0])
dones       = torch.tensor([0.0, 0.0, 1.0])
terminated  = torch.tensor([0.0, 0.0, 1.0])
gamma, lam = 0.5, 0.5

adv = torch.zeros_like(rewards)
running = 0.0
for t in reversed(range(len(rewards))):
    bootstrap_mask = 1.0 - terminated[t]
    continuation_mask = 1.0 - dones[t]
    delta = rewards[t] + gamma * bootstrap_mask * next_values[t] - values[t]
    running = delta + gamma * lam * continuation_mask * running
    adv[t] = running
print("advantages", adv.tolist())
print("returns   ", (adv + values).tolist())

advantages [1.375, 1.5, 0.0]
returns    [2.375, 3.5, 3.0]

# A column baseline broadcasts against a vector into an
# unintended pairwise matrix.
adv = torch.tensor([1.375, 1.5, 0.0])
values_column = torch.tensor([[1.0], [2.0], [3.0]])
print(tuple((adv + values_column).shape))

(3, 3)

Application type mismatch: The argument
  rlValTwo
has type
  Tensor Float [2]
but is expected to have type
  Tensor Float [3]
in the application
  Spec.RL.returnsFromAdvantages rlAdvThree rlValTwo

The PPO Objective🔗

For the sampled action at timestep t, the importance ratio is

r_t(\theta) =\exp\!\left( \log\pi_\theta(a_t\mid s_t) -\log\pi_{\theta_{\mathrm{old}}}(a_t\mid s_t) \right),

and the clipped surrogate that Schulman et al. (2017) maximize is

L^{\mathrm{clip}}_t(\theta) =\min\!\left( r_t(\theta)A_t,\; \operatorname{clip}(r_t(\theta),1-\epsilon,1+\epsilon)A_t \right).

The minimum makes clipping depend on the advantage sign. Take \epsilon=1/4 and a positive advantage A_t=1. At ratio one, the objective is one:

1.000000

At ratio 3/2, the clipped branch reaches 1+\epsilon:

1.250000

1.25 rather than 1.5: the gradient of this term with respect to the ratio is zero, removing this sample's incentive to increase the ratio further. This does not enforce a hard bound on policy movement: other samples, shared parameters, and optimizer state still affect it. Pull the ratio down to 1/2 and nothing is clipped, because the minimum picks the unclipped branch:

0.500000

With a negative advantage, multiplication by A_t reverses the order of the two candidate terms:

-1.500000

-0.750000

At ratio 3/2 with A_t=-1, the objective is -1.5, since \min(-1.5,-1.25)=-1.5. The clipped objective keeps penalizing an increase in this ratio beyond 1+\epsilon; it does not flatten there. In the opposite direction it flattens at -0.75 once the ratio falls below 1-\epsilon. Thus the objective is not a symmetric constraint on policy movement. For a categorical action with fixed positive old probability, the ratio also has the separate probability bound r_t\le 1/\pi_{\mathrm{old}}(a_t\mid s_t).

The PyTorch expression gives the same five numbers:

# Evaluate both signs with one shared clipping interval.
ratio = torch.tensor([1.0, 1.5, 0.5, 1.5, 0.5])
advan = torch.tensor([1.0, 1.0, 1.0, -1.0, -1.0])
eps = 0.25
obj = torch.min(ratio * advan,
                torch.clamp(ratio, 1 - eps, 1 + eps) * advan)
print(obj.tolist())

[1.0, 1.25, 0.5, -1.5, -0.75]

The full loss to minimize adds a critic term and an optional entropy bonus,

-L^{\mathrm{clip}}_t +c_v\,\bigl(V_\theta(s_t)-R_t\bigr)^2 -c_e\,H\!\left[\pi_\theta(\cdot\mid s_t)\right],

and the differentiable batch version over backend references lives in NN.Runtime.RL.PolicyGradient.Autograd, where the categorical log probability is built from logSoftmax and a one-hot action tensor. The pure scalar functions above are the scalar formulas without a tape. The ratio-based clipped objective can be checked exactly at ℚ; computing the ratio from log probabilities additionally requires an exponential. Their tape implementations live inside the autograd runtime and need their own agreement arguments.

The ratio is computed for the action actually stored in the rollout, not for whichever action the updated actor now prefers. A ratio above one says that sampled action has become more likely under the new policy; the advantage supplies the judgment about whether that change is desirable. The five examples separate those two pieces by providing ratios directly. They do not evaluate log probabilities or sample an action, which is why they can isolate the clipped minimum without involving a categorical distribution or a random seed.

A training loss minimizes the negative surrogate. At positive advantage, increasing the ratio inside the unclipped interval improves the surrogate and lowers that loss. The critic term has a different purpose: it fits the return target used to estimate advantages. The entropy term rewards a less concentrated action distribution when its coefficient is positive. These terms can pull shared parameters in different directions, so the flat branch of one sample's surrogate is not a bound on the complete update.

Checked Arithmetic For Returns And Ratios🔗

Returns, advantages, and importance ratios give us three places to look for non-finite values. A return accumulates over a horizon, a ratio is an exponential of a difference of logarithms, and normalization divides by a standard deviation that can be zero. Binary32 arithmetic can produce an infinity or a NaN and continue the calculation with that value (Goldberg, 1991).

The checked helpers under NN.Runtime.RL.Numerics.Float32 run the same recurrences over an executable binary32 semantics and return Except String .... They are described in the float32 soundness chapter; here we use them.

Start at the boundary. Casting a host binary64 value into the executable binary32 type can overflow, and the cast refuses rather than quietly returning an infinity:

"rejected before it reached the update"

The full error message names the input and the resulting value. I match on the constructor here to keep the rejection of 10^{39} readable. An accepted value carries its bit pattern, which is a reminder that this type is a binary32 encoding rather than a real number:

Except.ok 1056964608

1056964608 is 0x3F000000, the binary32 encoding of 0.5.

Now the discounted backup r+\gamma(1-d)\,V(s'), first on ordinary inputs:

Except.ok 2.000000

Now supply a bootstrap value that is already infinite:

Except.error "RL float32: non-finite configured binary32 value at discountedBackup/mul(t1,bootstrap): inf"

The label identifies the multiplication that produced the non-finite value: \gamma\cdot(1-d)\cdot V(s'). Larger checked routines retain the labels of their primitive operations, so the diagnostic locates the failing calculation within the update.

For background on formal floating-point analysis, see (Boldo and Melquiond, 2011). Here we can inspect the same recurrence with an interval calculation: compute the returns in binary32 and propagate intervals alongside them, then check whether each binary32 result lies inside its interval:

Except.ok true

The successful interval check returns Except.ok true: the outer constructor says the checked binary32 calculations completed, and the Boolean says their returned values passed the enclosure test. Those are separate conditions. A failed cast or intermediate overflow would return an error before there were finite results to compare. The dyadic three-reward example is chosen so its expected values are already known; it demonstrates how the numerical checks connect to the recurrence rather than estimating error for an unknown long trajectory.

discountedBackup_eq_ok : ∀ (reward gamma bootstrap : Float32Exec) (done : Bool) (out : Float32Exec),
  discountedBackupChecked reward gamma bootstrap done = Except.ok out →
    ExecFloat.Binary.isFinite (ExecFloat.mul gamma (continueMask done)) = true ∧
      ExecFloat.Binary.isFinite ((ExecFloat.mul gamma (continueMask done)).mul bootstrap) = true ∧
        ExecFloat.Binary.isFinite (ExecFloat.add reward ((ExecFloat.mul gamma (continueMask done)).mul bootstrap)) =
            true ∧
          out = discountedBackup reward gamma bootstrap done

Bellman Operators And Fixed Points🔗

The same vocabulary supports the classical theory. For a fixed policy \pi,

(T^\pi V)(s) =r(s,\pi(s)) +\gamma\mathbb E_{s'\sim P(\cdot\mid s,\pi(s))}[V(s')],

and the optimality operator takes the best action instead:

(T^\star V)(s) =\max_a\left( r(s,a)+\gamma\mathbb E_{s'\sim P(\cdot\mid s,a)}[V(s')] \right).

For 0\le\gamma<1 both are contractions in the sup metric, which is the single fact from which value iteration, uniqueness of the fixed point, and the familiar geometric error bounds follow. TorchLean proves it for finite stochastic MDPs:

@bellmanOptimality_contraction : ∀ {nStates nActions : ℕ} [inst : Fact (0 < nStates)] [inst_1 : Fact (0 < nActions)]
  (mdp : MDP nStates nActions),
  Valid mdp →
    ∀ (values₁ values₂ : Spec.RL.ValueFunction ℝ nStates),
      valueSupDist (bellmanOptimality mdp values₁) (bellmanOptimality mdp values₂) ≤
        mdp.discount * valueSupDist values₁ values₂

The discount bound arrives through Valid mdp, which also requires the transition rows to be probability distributions; the state and action counts are positive by instance arguments; and valueSupDist is a Finset.sup' over states, which is why nonemptiness has to be available. Nothing here is inferred from a simulator run.

Iterating that inequality gives a geometric error bound for value iteration:

@bellmanOptimality_iterate_error_to_fixedPoint : ∀ {nStates nActions : ℕ} [inst : Fact (0 < nStates)]
  [inst_1 : Fact (0 < nActions)] (mdp : MDP nStates nActions),
  Valid mdp →
    ∀ (v vStar : Spec.RL.ValueFunction ℝ nStates),
      bellmanOptimality mdp vStar = vStar →
        ∀ (k : ℕ), valueSupDist ((bellmanOptimality mdp)^[k] v) vStar ≤ mdp.discount ^ k * valueSupDist v vStar

k applications of the operator shrink the distance to a fixed point by \gamma^k. The companion theorems prove that such a fixed point is unique for both operators, and related results appear for deterministic and stochastic settings: first for deterministic finite MDPs in NN.Proofs.RL.MDP, once for stochastic ones in FiniteStochasticMDP, and once more for general measurable state spaces in NN.Proofs.RL.MarkovMDP.

The sup distance asks for the largest error over all states. A stochastic transition row averages successor errors using nonnegative weights that sum to one, so this averaging cannot exceed that largest error. Discounting then multiplies it by \gamma, which is the source of the contraction factor. This explains both parts of Valid: probability-row conditions control the averaging, and a discount below one makes the factor contractive. Neither condition concerns how a neural network is initialized or optimized.

The iteration counter in this theorem counts exact Bellman applications. A PPO epoch instead updates policy and value-network parameters using a sampled surrogate objective, so its logged epoch number cannot be substituted for k in the displayed bound. To use the bound for an approximate value-iteration implementation, one would also need to quantify the error introduced at each approximate application. The fixed-point hypothesis supplies the comparison function vStar; the theorem then controls distance to it from any initial value function.

@toFiniteStochasticMDP_valid : ∀ {width height : ℕ} (gw : Spec.RL.Envs.GridWorld width height),
  0 ≤ gw.discount → gw.discount < 1 → Valid gw.toFiniteStochasticMDP

External Environments🔗

GridWorld's transition function is a Lean definition that proofs can refer to. CartPole instead runs in Gymnasium, a Python library. Its transitions enter Lean through a checked communication boundary.

The bridge uses a subprocess speaking JSON lines, one object per line, and checks a contract on the Lean side of the exchange. The server is scripts/rl/gymnasium_server.py, and its docstring explains that you can run the file directly and type the protocol by hand. The exchange below shows what crosses the boundary:

$ python3 scripts/rl/gymnasium_server.py --env-id CartPole-v1
{"cmd": "describe"}
{"ok":true,"n_actions":2,"obs_shape":[4],"obs_dtype":"float32"}
{"cmd": "reset", "seed": 0}
{"ok":true,"obs":[0.013696168549358845,
                  -0.023021329194307327,
                  -0.04590264707803726,
                  -0.04834723472595215]}
{"cmd": "step", "action": 0}
{"ok":true,"obs":[0.013235742226243019,
                  0.17272774875164032,
                  -0.04686959087848663,
                  -0.3551521897315979],
 "reward":1.0,"terminated":false,"truncated":false}
{"cmd": "step", "action": 99}
{"ok":false,"error":"AssertionError: 99 (<class 'int'>) invalid"}
{"cmd": "close"}
{"ok":true}

The observation lines are wrapped here to fit the page; on the wire each reply is one line. The describe reply supplies the dimensions Lean uses to construct its tensor shapes. Gymnasium's terminated and truncated arrive as separate flags: an external time-limit truncation generally retains the value bootstrap, while termination ends it. The PPO collector records both the done episode boundary and the task's terminated flag. GAE uses terminated to mask the next-state value and done to stop advantage continuation. A truncated step therefore bootstraps from the final observation before auto-reset, without including rewards from the following episode. The rollout viewer uses the same calculation as training. The generic single-mask GAE API retains its existing semantics.

Python refuses the out-of-range action. The close command gets an acknowledgement, so the Lean side can distinguish a clean shutdown from a crash.

The describe handshake fixes the observation and action contract before training tensors are constructed. CartPole supplies four real-valued observations and two actions, while Pong RAM supplies byte-like observations and a different action space. Once that boundary has been checked, the common collector can use the same rollout structure. The Python process still owns the transition dynamics; successful JSON exchange establishes that the data reached Lean in the expected form, rather than proving how the simulator produced it.

The rejected action 99 probes one boundary without changing the observation dimensions. It shows why a successful shape handshake is only the beginning of the protocol: every subsequent message still carries values that need checking. Conversely, a valid action does not imply that its resulting reward is finite or its observations lie in the declared range. The following contracts separate these conditions, so a failure can identify the field responsible instead of reporting only that an episode could not continue.

Except.ok 0.950000

Except.error "RL boundary: expected finite reward, got NaN."

Except.error "RL boundary: reward out of range: got 7.500000, expected in [-1.050000, 1.000000]."

Except.error "RL boundary: nextObservation tensor out of range: expected all entries in [0.000000, 1.000000]."

Except.error "RL boundary: both `terminated` and `truncated` are true (contract requires exclusivity)."

Except.error "RL boundary: action out of range: 7 (nActions=4)."

@contractHolds_of_checkTransitionFin_eq_ok : ∀ {obsShape : Shape} {nActions : ℕ} (c : Contract obsShape nActions)
  (observation nextObservation : Tensor Float obsShape) (action : Fin nActions) (reward : Float)
  (terminated truncated : Bool) (t : Transition obsShape nActions),
  checkTransitionFin c observation nextObservation action reward terminated truncated = Except.ok t → ContractHolds c t

CartPole And Pong RAM Results🔗

CartPole and Pong RAM exercise the same collector with external environments.

ppo_cartpole runs the same PPO loop against Gymnasium. Here is a recorded 30-update run with the shipped hyperparameters, evaluating every fifth update over five episodes with a 200-step cap:

$ torchlean ppo_cartpole --updates 30 --eval-every 5 \
    --eval-episodes 5 --eval-max-steps 200 \
    --log-json /mnt/build/rl-cartpole-log.json
[TorchLean] arithmetic: native binary32
[TorchLean] execution: eager
[TorchLean] device: cpu
ppo_cartpole: PPO on CartPole-v1 (horizon=64) (device=cpu)
  env: Python Gymnasium subprocess (JSON-lines bridge)
       + Lean boundary contract
  eval(step=0) avg_return=9.600000
  update=5 avg_return=10.000000
  update=10 avg_return=9.800000
  update=15 avg_return=9.800000
  update=20 avg_return=9.800000
  update=25 avg_return=10.000000
  update=30 avg_return=10.000000
  wrote TrainLog JSON: /mnt/build/rl-cartpole-log.json
ppo_cartpole: done
ppo_cartpole: ok

The evaluated greedy policy survives about ten steps both before and after training. With horizon 64, thirty updates collect roughly two thousand environment steps. The collector and updates completed, but the measured return shows little improvement. That observation alone does not tell us whether collecting more steps would improve this policy.

For CartPole, avg_return is capped by --eval-max-steps: an episode that survives all eight steps under an eight-step cap reports 8.000000. That measurement cannot distinguish policies that survive eight steps from those that survive much longer. The run above uses 200, and the JSON log records the cap alongside the curve.

ppo_pong_ram is the third example, and requires the external ALE package and environment assets. An environment missing that package reports the following failure:

$ torchlean ppo_pong_ram --updates 1
[TorchLean] arithmetic: native binary32
[TorchLean] execution: eager
[TorchLean] device: cpu
ppo_pong_ram: PPO on ALE/Pong-v5
  (obs=ram, horizon=128) (device=cpu)
  env: Python Gymnasium subprocess (ALE)
       + Lean boundary contract
  starting env: ALE/Pong-v5 (obs_type=ram)
Traceback (most recent call last):
  ...
ModuleNotFoundError: No module named 'ale_py'
  ...
RuntimeError: Requested an `ALE/...` environment id, but
`ale-py` could not be imported/registered. Detected
gymnasium=1.3.0, ale-py=not installed. Try: python3 -m pip
install --user --upgrade ale-py 'gymnasium>=1.0'
error: Gymnasium: unexpected EOF from server
  (process exited?)

Two tracebacks are elided. The RuntimeError names the missing package, reports the installed gymnasium version, and gives the command to fix it. The Lean-side exception on the last line says only that the pipe closed. The diagnostic is visible because the subprocess writes to the same standard error, not because the Lean side captured it, so in a context that separates the two streams the useful half can go missing. Keep both output streams when diagnosing an external environment failure.

The CartPole trace illustrates why evaluation settings belong beside the metric. Thirty update cycles can finish successfully while the greedy policy's survival time remains near ten steps. A different exploration distribution might collect different trajectories even when greedy evaluation looks unchanged. The trace reports what this evaluation observed; it does not isolate whether the limiting factor is optimization, advantage estimates, rollout coverage, or model capacity. The Pong failure occurs earlier still, before the environment can provide a trajectory.

Off-Policy Data: Replay Buffers And DQN🔗

PPO reuses a collected rollout for its optimization epochs, then collects another rollout. DQN retains past transitions for later updates (Mnih et al., 2015). Drawing random minibatches from that buffer lets an update combine transitions collected at different times. The small example below uses a deterministic sampler whose selected entries we can inspect.

A replay buffer must keep its size within its capacity. Pushing three transitions into a buffer of capacity two demonstrates the full-buffer case:

(2, #[2.000000, 3.000000])

The oldest transition is gone and the size stayed at the capacity. A bounded Python deque has the same eviction behavior:

# Compare oldest-first eviction with a capacity-two Python
# container.
>>> from collections import deque
>>> b = deque(maxlen=2)
>>> for k in (1.0, 2.0, 3.0): b.append(k)
>>> list(b)
[2.0, 3.0]

TorchLean also proves the size after pushing into a full buffer:

@push_size_of_full : ∀ {α : Type} {obsShape : Shape} {nActions : ℕ} (b : Buffer α obsShape nActions)
  (t : Transition α obsShape nActions), 0 < b.capacity → b.items.size = b.capacity → (b.push t).items.size = b.capacity

push_size_of_room says the size increases by one when it is below capacity; push_size_of_full says it stays at capacity when the buffer is full and capacity is positive. These are size statements under their respective hypotheses. The push definition appends the new transition and, if necessary, removes the oldest entry, but the displayed size theorem alone does not establish which transitions were retained.

The type also pins the transition down. Transition Float [1] 2 fixes the observation shape and the action count, and the action field is a Fin 2, which is why the literal above needs by decide. A deque does not check these fields when an item is appended; its callers must arrange those checks separately.

The buffer stores oldest-first entries in a bounded FIFO array. Its capacity controls how much past experience remains available, while the sampling policy controls which retained entries participate in an update. These are separate choices. The capacity theorem assumes an initially full valid buffer and positive capacity; a manually constructed oversized buffer is not covered by that hypothesis. The public empty constructor and repeated pushes are the path illustrated here, so the capacity-two example exposes eviction without requiring a large training run.

$ torchlean dqn_replay
dqn_replay: begin
stored transitions: 2
sampled transitions: 4
DQN minibatch MSE loss:   0.917800
DQN minibatch Huber loss: 0.452500
soft target update example: 1.000000
dqn_replay: ok

([0.200000, 2.000000, 0.500000], [0.100000, 2.400000, 0.300000])

1.345600

0.490000

0.660000

0.245000

0.917800

0.452500

6.000000

1.000000

1.040000

softUpdateScalar_sub_target_real : ∀ (tau online target : ℝ),
  softUpdateScalar tau online target - target = tau * (online - target)

Implementation And Proof Coverage🔗

The source directories separate specifications, runtime code, and proofs; the command registry identifies maintained examples. The table pairs each implementation area with its proof scope.

The MDP theorems apply to their named mathematical environments, including the sparse-reward GridWorld. The executable PPO example uses the shaped reward described above. For a deployment, check the selected environment contract, termination policy, arithmetic, and runtime; the earlier sections state these contracts beside the corresponding code.

Contract Exercises🔗

1. Run torchlean ppo_gridworld --updates 40 --eval-episodes 1 --eval-max-steps 8 --log /tmp/rl.json --path /tmp/rl-path.json. Recompute the final return from that path and the shaped reward, rather than assuming it follows the one-update path shown above.

2. Change the discount in the rlGrid definition above to 1 and see which hypothesis of toFiniteStochasticMDP_valid fails, and what the error message says.

3. Evaluate discountedReturns at ℚ on a reward tensor of your own, then evaluate the same tensor at Float, and find the shortest trajectory where the two answers disagree once you convert.

4. Write the GAE recursion in NumPy for \lambda=0.5 and cross-check three horizons against generalizedAdvantageEstimation. Then break the shapes on purpose and compare the two error messages.

5. Set requireExclusiveDoneFlags := false in rlContract and rerun the both-flags-true check. Then decide which setting you would want for a Gymnasium wrapper you did not write.

6. Change batchSize := 4 to batchSize := 16 in the DQN replay example source and predict the repeated sample loss before rebuilding and running it; the command has no batch-size flag.

7. Feed ofFloatChecked the largest binary64 value that still converts, then the next one up, and read both messages.

8. Run ppo_cartpole with --eval-max-steps 8 and then with --eval-max-steps 500, and compare how the cap limits what each return curve can establish.

Related Verification Chapters🔗

Verification Overview explains how executable checks connect to mathematical claims. Scientific ML Verification examines the corresponding obligations for PDE surrogates, including the meaning of a measured residual.