7.4. Reinforcement Learning
In supervised learning, the dataset is usually fixed before training starts. In reinforcement learning, the current policy helps create its own future data. A complete application therefore contains more than a neural network:
\text{environment}
\longrightarrow\text{transition}
\longrightarrow\text{rollout or replay}
\longrightarrow\text{return and advantage}
\longrightarrow\text{policy/value update}.
Each arrow carries assumptions. Is the observation shape correct? Is the action valid? Is the
reward finite? Does done mean termination, truncation, or either? Did the rollout keep its fields
aligned? TorchLean gives those questions explicit runtime and specification objects.
There are three complementary layers:
-
NN.Spec.RLdefines environments, MDPs, Bellman operators, returns, and advantages; -
NN.Runtime.RLimplements checked transitions, replay, rollouts, PPO, and Gymnasium communication; -
NN.Proofs.RLproves structural, numerical, and dynamic-programming facts about named objects from the first two layers.
The easiest way to see the pieces together is a small GridWorld run.
7.4.1. A Complete PPO GridWorld Run
The environment is a 4\times4 GridWorld defined in Lean. The actor and critic use a fixed rollout
horizon of 64. Run one update on CPU and send the artifacts to temporary files:
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
The current checkout produces:
[TorchLean] dtype: Float (Lean `Float`, trusted runtime semantics) [TorchLean] backend: Runtime.Autograd.Torch.Backend.eager [TorchLean] device: cpu torchlean 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=0 avg_return=3.600000 wrote TrainLog JSON: /tmp/ppo-gridworld-trainlog.json torchlean ppo_gridworld: wrote policy snapshot to /tmp/ppo-gridworld-policy.json torchlean ppo_gridworld: wrote path snapshot to /tmp/ppo-gridworld-path.json torchlean ppo_gridworld: done torchlean ppo_gridworld: ok
This trace contains three different results:
-
the PPO/autograd program completed one update;
-
a greedy-policy evaluation changed on this seeded run;
-
the command wrote a scalar curve, a policy table, and a decoded path.
Only the first is an optimizer execution fact. The second is an empirical observation from one small run. The third gives inspectable evidence. The formal MDP facts described below are separate theorems.
The implementation is
NN/Examples/Models/RL/PPOGridWorld.lean.
The pure environment is
NN.Spec.RL.Envs.GridWorld.
7.4.2. One Actor, Two Views Of A Batch
The GridWorld command does not hide its actor and critic inside the training loop. It starts with a small reusable configuration:
def modelCfg : nn.models.PPOActorCriticConfig :=
{ obsDim := 16, hiddenDim := 32, nActions := 4 }
NN.API.Models.PPO
then builds a two-layer tanh MLP for each role. The actor ends in four action logits; the critic
ends in one value. Both constructors preserve a caller-supplied prefix shape, so the same model
description has two useful views:
Prefix | Actor shape | Critic shape |
|---|---|---|
one observation |
|
|
rollout of 64 observations |
|
|
The prefix belongs to the sequential model's input and output shapes rather than an implicit broadcast convention. GridWorld, CartPole, and Pong RAM reuse these constructors with different observation and action dimensions; environment collection and boundary checks remain outside the model helper. This separation is why the same actor definition can score one state during rollout and a full horizon during optimization without introducing a second, loosely related network.
7.4.3. What Enters A Rollout
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,\;V(s_t),\;V(s_{t+1})).
The Lean structure in
NN.Runtime.RL.PPO.Rollout
uses:
-
state : Tensor α obsShape; -
action : Fin nActions; -
scalar old log probability, reward, value, and next value;
-
a Boolean episode-boundary marker.
A Rollout α obsShape nActions horizon contains an array plus a proof that its size is exactly
horizon. This avoids the “parallel arrays drifted out of sync” failure common in hand-written
buffers.
Converting a rollout to an actor-critic minibatch produces:
\begin{aligned}
\text{states}&:T\times\text{obsShape},\\
\text{actionsOneHot}&:T\times A,\\
\text{oldLogProb}&:T,\\
\text{advantages}&:T,\\
\text{valueTarget}&:T\times1.
\end{aligned}
The conversion is total because the horizon invariant is already present.
7.4.4. Returns And Generalized Advantage Estimation
Let d_t be one when timestep t ends an episode and zero otherwise. The one-step temporal
difference residual is
\delta_t
=r_t+\gamma(1-d_t)V(s_{t+1})-V(s_t).
Generalized Advantage Estimation runs backward:
A_t
=\delta_t+\gamma\lambda(1-d_t)A_{t+1}.
The corresponding value target is
R_t=A_t+V(s_t).
TorchLean keeps the list-level definitions in
NN.Spec.RL.Core
and tensor-shaped runtime siblings in
NN.Runtime.RL.Core.
Rollout.toActorCriticSample computes value targets from the unnormalized advantages, then
z-score normalizes advantages for the policy term.
This is more than a documentation choice. Months after an experiment, the definition answers questions that a log cannot: whether truncation stopped bootstrapping, whether normalization changed value targets, and which direction the recurrence traversed the rollout.
7.4.5. The PPO Objective
For the sampled action at timestep t, define
r_t(\theta)
=\exp\!\left(
\log\pi_\theta(a_t\mid s_t)
-\log\pi_{\theta_{\mathrm{old}}}(a_t\mid s_t)
\right).
The clipped surrogate 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 actor minimizes the negative mean surrogate. The critic receives a value loss, and the full
autograd program can add an entropy term. The categorical action log probability is implemented
from logSoftmax and a one-hot action tensor in
NN.Runtime.RL.PolicyGradient.Autograd.
The ratio, clipping, return, and advantage computations also have checked binary32 helpers under
NN.Runtime.RL.Numerics.Float32.
They use the executable IEEE32Exec semantics and return Except String ..., so a NaN, infinity,
or failed finite-path precondition is visible rather than entering the update silently.
These helpers establish properties of selected scalar recurrences; whole-loop CUDA refinement is a larger claim.
The corresponding proof bridge is split between
Floats.IEEE32Exec
and
Floats.CheckedRuntime.
The first states finite-path equalities for discounted backups and TD residuals; the second uses
the checked wrappers to discharge those finite-result premises. Together they cover those selected
recurrences, not a numerical error theorem for the whole PPO loop.
7.4.6. Bellman Operators And Fixed Points
The same vocabulary supports classical MDP theorems. 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')].
The optimality operator is
(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\leq\gamma<1, both are contractions in the sup norm:
\|TV-TW\|_\infty
\le\gamma\|V-W\|_\infty.
The Markov-MDP proof file
NN.Proofs.RL.MarkovMDP
contains the named results bellmanPolicy_contraction,
bellmanOptimality_contraction, bellmanPolicy_fixedPoint_unique, and
bellmanOptimality_fixedPoint_unique.
These are real dynamic-programming theorems. PPO convergence concerns a different object: a parameterized stochastic policy updated from sampled finite trajectories.
For explicitly finite state and action spaces,
FiniteStochasticMDP
also proves monotonicity, contraction, uniqueness of policy and optimality fixed points, and
geometric iterate-error bounds. The finiteness and discount hypotheses are part of those theorem
statements; they are not inferred from a simulator run.
7.4.7. External Environments
GridWorld is fully represented in Lean. CartPole and Pong are not. Reimplementing every simulator inside the theorem prover would make TorchLean isolated from the ecosystem, so the runtime can start a Python Gymnasium process and communicate through one JSON request and response per line.
The bridge lives in
NN.Runtime.RL.Gymnasium.
At startup, Lean asks the process for the observation shape and number of actions. On reset and
step, the returned values pass through
NN.Runtime.RL.Boundary.
A boundary contract may require:
-
every observation entry to be finite;
-
every reward to be finite;
-
observations and rewards to lie in declared closed intervals;
-
actions to satisfy
action < nActions; -
terminatedandtruncatednot to be simultaneously true.
After checking, an action is a Fin nActions and the transition is a
Spec.RL.ObservedTransition. A malformed Python dictionary never becomes training data merely
because it could be decoded as JSON.
This is a checked interface, not a proof of the simulator. TorchLean does not establish that Gymnasium's CartPole dynamics are Markov, that ALE emulates the Atari hardware correctly, or that the Python process honored its random seed.
7.4.8. CartPole And Pong RAM
Install the optional external dependency:
python3 -m pip install --user 'gymnasium>=1.0'
Then a short CartPole run is:
lake exe torchlean ppo_cartpole --device cpu \ --updates 1 --eval-every 1 \ --eval-episodes 1 --eval-max-steps 8
The command expects a four-entry observation and two actions. Its actor and critic are MLPs with
hidden width 32, rollout horizon 64, discount 0.99, GAE parameter 0.95, and Adam learning rate
3\times10^{-4}. Those values are defined in
PPOCartPole.lean,
not inferred from a generic PPO label.
The optional Pong path uses RAM observations to keep the JSON-lines boundary inspectable:
python3 -m pip install --user 'gymnasium>=1.0' ale-py lake exe torchlean ppo_pong_ram --device cpu \ --check-env-only
The environment is ALE/Pong-v5, each observation has 128 byte-valued entries, and the action
space has six elements. A pixel PPO implementation would require a much higher-throughput
transport; the RAM example is intentionally about the checked environment boundary.
7.4.9. Off-Policy Data: DQN Replay
Off-policy algorithms reuse transitions. TorchLean's
Replay.Buffer
is a bounded FIFO array indexed by scalar type, observation shape, and action count. Pushing into a
full buffer drops the oldest transition; sampling is deterministic from explicit seed and counter
values.
Run the self-contained example:
lake exe torchlean dqn_replay
Current output:
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
The sample count can exceed the stored count because this compact demonstration samples with replacement. The example computes DQN targets, MSE and Huber losses, and a soft target-network update. It is not a full epsilon-greedy environment training run.
The proof module
NN.Proofs.RL.Replay
establishes structural facts about empty buffers, zero capacity, size growth, and eviction at
capacity. Those are small theorems with a useful job: later DQN reasoning need not assume that the
storage layer preserved its own capacity invariant.
The
DQN algebra proof
adds exact real identities for target-network soft updates, including the zero/one mixing endpoints
and softUpdateScalar_sub_target_real. It proves the scalar recurrence, not correctness of replay
sampling, Q-network optimization, or a complete DQN agent.
7.4.10. Hands-On Checks
7.4.10.1. Inspect The Artifacts
After the GridWorld command, compare the before and after policy arrays:
python3 -m json.tool /tmp/ppo-gridworld-policy.json python3 -m json.tool /tmp/ppo-gridworld-path.json
The policy artifact records one greedy action per grid cell. The path artifact records decoded
(row, column) states. A rising scalar return is easier to interpret when these two objects agree
with it.
7.4.10.2. Exercise Strict Parsing
The application runner rejects unconsumed flags. For example, --rollout 8 is not a GridWorld
option:
lake exe torchlean ppo_gridworld --device cpu --updates 1 --rollout 8
The command fails with an unexpected arguments message instead of silently running horizon 64.
The horizon is a typed constant in this example; changing it is a source-level model change.
7.4.10.3. Compare CPU And CUDA
With a CUDA build:
lake -R -K cuda=true exe torchlean ppo_gridworld --device cuda \ --updates 1 --eval-every 1 --eval-episodes 1 --eval-max-steps 8 \ --show-backend
The environment dynamics remain Lean-native. The actor, critic, and PPO autograd operations move to
the selected CUDA runtime, and --show-backend prints the chosen backend contracts. A successful
CUDA run is runtime evidence; kernel-level proof status is read from those contracts rather than
inferred from the device name.
7.4.11. What Is Implemented Today
TorchLean currently provides:
-
executable PPO applications over a Lean environment and external Gymnasium environments;
-
typed fixed-horizon rollouts and bounded replay;
-
explicit checks for external observations, actions, rewards, and episode flags;
-
checked binary32 helpers for selected return, advantage, and PPO scalar formulas;
-
Bellman contraction and fixed-point theorems for named MDP objects;
-
finite-state policy/optimality iteration error bounds;
-
finite-path IEEE32Exec and checked-runtime bridges for selected backup and TD formulas;
-
real-algebra target-network update laws for DQN;
-
structural proofs for environment and replay components.
The Bellman results cover the named MDP objects, while the PPO commands exercise the training system. Connecting a particular simulator and native PPO run all the way to those mathematical objects remains a larger end-to-end result.
7.4.12. References
-
Sutton and Barto, Reinforcement Learning: An Introduction, second edition.
-
Schulman et al., High-Dimensional Continuous Control Using Generalized Advantage Estimation, 2015.
-
Schulman et al., Proximal Policy Optimization Algorithms, 2017.
-
Mnih et al., Human-level control through deep reinforcement learning, 2015.
-
Gymnasium, environment API.