Implementation:LaurentMazare Tch rs PPO Agent
| Knowledge Sources | |
|---|---|
| Domains | Reinforcement Learning, Deep Learning, Game AI |
| Last Updated | 2026-02-08 00:00 GMT |
Overview
Implements Proximal Policy Optimization (PPO) for training an agent to play Atari games (SpaceInvaders) using a convolutional actor-critic neural network.
Description
This module implements the PPO algorithm as described in Schulman et al. 2017 ("Proximal Policy Optimization Algorithms"). The implementation uses a shared convolutional backbone with separate actor (policy) and critic (value) heads. The convolutional network processes stacked frames (4 consecutive 84x84 grayscale frames) through three convolutional layers followed by a fully connected layer.
The training loop collects rollouts across 8 parallel environments for 256 steps each, then computes discounted returns with a discount factor of 0.99. Optimization is performed using mini-batch updates over 4 epochs with a batch size of 64, employing the Adam optimizer with a learning rate of 1e-4 and gradient clipping at 0.5.
A FrameStack utility manages the observation buffer, stacking the most recent 4 frames and zeroing out frames when episodes terminate (using done masks). The loss function combines value loss (MSE of advantages, weighted by 0.5), policy loss (negative advantage-weighted log-probabilities), and an entropy bonus (weighted by 0.01) to encourage exploration.
The module also provides a sample function for evaluating a trained model by loading saved weights and running inference without gradient computation.
Usage
Use this implementation when training a reinforcement learning agent on Atari environments using PPO. It requires the cpython crate for interacting with OpenAI Gym via the VecGymEnv wrapper. The trained weights are periodically saved in .ot format every 1000 updates.
Code Reference
Source Location
- Repository: LaurentMazare_Tch_rs
- File: examples/reinforcement-learning/ppo.rs
- Lines: 1-181
Signature
type Model = Box<dyn Fn(&Tensor) -> (Tensor, Tensor)>;
fn model(p: &nn::Path, nact: i64) -> Model
struct FrameStack {
data: Tensor,
nprocs: i64,
nstack: i64,
}
impl FrameStack {
fn new(nprocs: i64, nstack: i64) -> FrameStack
fn update<'a>(&'a mut self, img: &Tensor, masks: Option<&Tensor>) -> &'a Tensor
}
pub fn train() -> cpython::PyResult<()>
pub fn sample<T: AsRef<std::path::Path>>(weight_file: T) -> cpython::PyResult<()>
Import
use super::vec_gym_env::VecGymEnv;
use tch::kind::{FLOAT_CPU, INT64_CPU};
use tch::{nn, nn::OptimizerConfig, Kind, Tensor};
I/O Contract
| Parameter | Type | Description |
|---|---|---|
| ENV_NAME | &str | Atari environment name ("SpaceInvadersNoFrameskip-v4") |
| NPROCS | i64 | Number of parallel environments (8) |
| NSTEPS | i64 | Steps per rollout (256) |
| NSTACK | i64 | Number of stacked frames (4) |
| UPDATES | i64 | Total training updates (1,000,000) |
| OPTIM_BATCHSIZE | i64 | Mini-batch size for optimization (64) |
| OPTIM_EPOCHS | i64 | Optimization epochs per update (4) |
| Output | Type | Description |
|---|---|---|
| Model output | (Tensor, Tensor) | Critic value and actor logits |
| Saved weights | .ot file | Model checkpoint saved every 1000 updates |
Usage Examples
// Training the PPO agent
// Requires Python environment with OpenAI Gym and atari_wrappers module
ppo::train().unwrap();
// Sampling from a trained agent
ppo::sample("trpo1000.ot").unwrap();