Skip to main content

thrust_rl/train/
ppo.rs

1//! PPO trainer (Burn backend).
2//!
3//! After phase 5 of the Burn migration (#82), Burn is the only tensor
4//! backend in the workspace. The historical `train::ppo_burn` parallel
5//! sibling has been collapsed back into `train::ppo`. The trainer struct
6//! retains its `PPOTrainerBurn<B, P, O>` name because Burn's
7//! optimizer-consumes-module ownership model is structurally different
8//! from a hypothetical in-place trainer.
9//!
10//! # Contents
11//!
12//! - [`actor_learner`] — single-host asynchronous actor-learner runner (N
13//!   inference-only actor threads feeding one learner over `crossbeam-channel`;
14//!   Phase 2 of the distributed-training epic #265).
15//! - [`config`] — `PPOConfig` hyperparameters / builder API.
16//! - [`stats`] — `TrainingStats` / `AggregatedStats` per-update metrics.
17//! - [`loss`] — backend-generic PPO loss math (policy/value/entropy) and the
18//!   `generate_minibatch_indices` helper.
19//! - [`trainer`] — `PPOTrainerBurn<B, P, O>` that owns the policy module
20//!   (Burn's optimizer-consumes-module ownership model) and exposes a
21//!   `train_step` that runs the surrogate-loss / gradient-step / KL early stop
22//!   logic.
23
24pub mod actor_learner;
25pub mod config;
26pub mod loss;
27pub mod recurrent_trainer;
28pub mod stats;
29pub mod trainer;
30
31pub use actor_learner::{
32    ActorChannels, ActorHandle, ActorStats, AsyncActorLearnerConfig, LearnerReport, actor_thread,
33    learner_loop, load_policy_from_broadcast, serialize_policy, spawn_actor,
34};
35pub use config::PPOConfig;
36pub use loss::{
37    compute_entropy_loss, compute_policy_loss, compute_value_loss, generate_minibatch_indices,
38    scalar_f64,
39};
40pub use recurrent_trainer::RecurrentPPOTrainer;
41pub use stats::{AggregatedStats, TrainingStats};
42pub use trainer::PPOTrainerBurn;