PPO on Atari Pong (RAM Observations) (Executable Example)

This example mirrors NN/Examples/Models/RL/PPOCartPole.lean, but targets an Atari game via the Arcade Learning Environment (ALE) registered into Gymnasium as ALE/Pong-v5.

Why "RAM" observations?

• Pixel-based Atari PPO is absolutely doable, but a JSON-lines subprocess bridge is not the right transport if you want millions of steps/hour. RAM observations (obs_type="ram", shape 128) keep the bridge compact and make this run viable as a native Lean executable.

The key TorchLean interface remains the same:

• Algorithm math (GAE, PPO clipped objective) is Lean definitions.
• Autograd program (PPO loss) is a TorchLean backend-generic program (CPU or CUDA).
• Trust boundary is explicit: every externally sampled transition is checked by Runtime.RL.Boundary.Contract before it can influence training.

Dependencies

Atari/ALE environments require ale-py and a recent gymnasium:

python3 -m pip install --user 'gymnasium>=1.0' ale-py

CLI flags

• --device 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>: write the widget log JSON to a custom path.

This command is optional in the sense that it depends on a compatible external ALE/Gymnasium installation. It is available through the runner but is not part of the default quick-check list.

Artifacts:

• Writes data/rl/ppo_pong_ram_trainlog.json by default (override with --log).
• Visualize it in the editor via NN/Examples/Models/RL/Views/PPOPongRam.lean.

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
• Machado et al., "Revisiting the Arcade Learning Environment: Evaluation Protocols and Open Problems" (2018): https://arxiv.org/abs/1709.06009
• ALE docs (environment catalogue and versioned ALE/...-v5 ids): https://ale.farama.org/
• Gymnasium API reference (reset/step, terminated vs truncated): https://gymnasium.farama.org/

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

Name used in CLI error messages and banners.

def NN.Examples.Models.RL.PPOPongRam.usage : String

Help text for the optional ALE/Pong RAM PPO runner.

Configuration

def NN.Examples.Models.RL.PPOPongRam.envId : String

Atari environment id passed to the Python subprocess.

def NN.Examples.Models.RL.PPOPongRam.gymServerScript : String

Relative path to the Python Gymnasium bridge script (spawned as a subprocess).

def NN.Examples.Models.RL.PPOPongRam.observationWidth : ℕ

Pong RAM observation dimension.

Gymnasium exposes RAM as Box(0, 255, (128,), uint8) when obs_type="ram".

def NN.Examples.Models.RL.PPOPongRam.actionCount : ℕ

Number of discrete actions in Pong under ALE's reduced action set.

def NN.Examples.Models.RL.PPOPongRam.hiddenWidth : ℕ

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

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

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

def NN.Examples.Models.RL.PPOPongRam.discountFactor : Float

Discount factor used in returns / GAE.

def NN.Examples.Models.RL.PPOPongRam.gaeLambda : Float

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

def NN.Examples.Models.RL.PPOPongRam.learningRate : Float

Adam learning rate used for the Pong RAM actor-critic update.

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

Number of PPO optimization epochs per collected rollout batch.

def NN.Examples.Models.RL.PPOPongRam.maxUpdates : ℕ

Default maximum number of PPO updates (override with --updates).

def NN.Examples.Models.RL.PPOPongRam.defaultEvaluationInterval : ℕ

Default evaluation checkpoint interval (override with --eval-every).

def NN.Examples.Models.RL.PPOPongRam.defaultEvaluationEpisodes : ℕ

Default evaluation episodes per checkpoint (override with --eval-episodes).

instance NN.Examples.Models.RL.PPOPongRam.instNeZeroNatHorizon : NeZero horizon

instance NN.Examples.Models.RL.PPOPongRam.instNeZeroNatActionCount : NeZero actionCount

def NN.Examples.Models.RL.PPOPongRam.observation : Spec.Shape

The observation tensor shape used by this run: [..., observationWidth].

def NN.Examples.Models.RL.PPOPongRam.rollout : Spec.Shape

def NN.Examples.Models.RL.PPOPongRam.rolloutStates : Spec.Shape

def NN.Examples.Models.RL.PPOPongRam.rolloutLogits : Spec.Shape

Logits for a whole rollout: one row of action logits per timestep.

def NN.Examples.Models.RL.PPOPongRam.rolloutValues : Spec.Shape

Value estimates for a whole rollout, one per timestep.

def NN.Examples.Models.RL.PPOPongRam.actionLogits : Spec.Shape

Action logits at a single timestep.

def NN.Examples.Models.RL.PPOPongRam.value : Spec.Shape

A single value estimate, shaped [1] so it composes with the batched shapes above.

Model (Actor + Critic)

We use MLPs over RAM. Pixel observations can instead use the arbitrary-rank convolutional model from TorchLean.nn.models.cnn after applying the appropriate Atari preprocessing.

@[reducible, inline]

abbrev NN.Examples.Models.RL.PPOPongRam.modelConfig : TorchLean.nn.models.PPO.Config

def NN.Examples.Models.RL.PPOPongRam.actor (batchShape : Spec.Shape := []) : TorchLean.nn.Builder (TorchLean.nn.Sequential (modelConfig.inputShape batchShape) (modelConfig.actorOutputShape batchShape))

Construct the actor network as an MLP mapping RAM observations to action logits.

def NN.Examples.Models.RL.PPOPongRam.critic (batchShape : Spec.Shape := []) : TorchLean.nn.Builder (TorchLean.nn.Sequential (modelConfig.inputShape batchShape) (modelConfig.criticOutputShape batchShape))

Construct the critic network as an MLP mapping RAM observations to a scalar value estimate.

Gymnasium / ALE bridge

We request RAM observations by passing {"obs_type": "ram"} to gym.make through the bridge's --make-kwargs option. The server also auto-registers ale_py when envId starts with ALE/.

def NN.Examples.Models.RL.PPOPongRam.makeKwargs : Array (String × Lean.Json)

def NN.Examples.Models.RL.PPOPongRam.contract : Runtime.RL.Boundary.Contract observation actionCount

What Lean insists on before it will believe anything the Python environment sends.

Every observation and reward must be finite, and RAM bytes must lie in [0, 255]. That last check is not about the emulator, which cannot produce anything else; it is about the protocol and the adapter between them, where a byte-order or dtype mistake would show up as out-of-range values.

def NN.Examples.Models.RL.PPOPongRam.checkEnvOnly : IO Unit

Start ALE, reset once, take one checked step, and close the subprocess.

This exercises the Gymnasium subprocess, ALE registration, RAM observation shape handshake, and Lean side boundary contract as the full PPO runner, without collecting a 128-step rollout.

Main Training Loop

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