Skip to main content

thrust_rl/train/
mod.rs

1//! Training algorithms (Burn backend).
2//!
3//! Hosts the PPO and DQN trainers, the backend-agnostic optimizer
4//! abstraction, and the loss math shared between them. After phase 5 of
5//! the Burn migration (#82), Burn is the only tensor backend in the
6//! workspace.
7
8/// Burn optimizer wrapper used by both PPO and DQN.
9pub mod optimizer;
10
11pub mod a2c;
12pub mod bc;
13pub mod dqn;
14pub mod ppo;
15pub mod sac;
16
17pub use a2c::{A2cConfig, A2cStats, A2cTrainer, compute_a2c_policy_loss, compute_a2c_value_loss};
18pub use bc::{BcConfig, BcEpochStats, BcTrainer, Demonstrations, compute_bc_loss};
19pub use dqn::{DQNConfig, DQNStepStatsBurn, DQNTrainerBurn};
20pub use optimizer::{BackendOptimizer, BurnOptimizer};
21pub use ppo::{
22    AggregatedStats, PPOConfig, PPOTrainerBurn, TrainingStats, compute_entropy_loss,
23    compute_policy_loss, compute_value_loss, generate_minibatch_indices,
24};
25pub use sac::{SacConfig, SacStepStats, SacTrainer};