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
11/// Global gradient-norm clipping shared by the PPO trainers and the joint
12/// multi-agent trainer (issues #239, #299).
13pub mod grad_clip;
14
15pub mod a2c;
16pub mod bc;
17pub mod dqn;
18pub mod ppo;
19pub mod sac;
20
21pub use a2c::{A2cConfig, A2cStats, A2cTrainer, compute_a2c_policy_loss, compute_a2c_value_loss};
22pub use bc::{BcConfig, BcEpochStats, BcTrainer, Demonstrations, compute_bc_loss};
23pub use dqn::{DQNConfig, DQNStepStatsBurn, DQNTrainerBurn};
24pub use grad_clip::{clip_grads_by_global_norm, global_grad_norm};
25pub use optimizer::{BackendOptimizer, BurnOptimizer};
26pub use ppo::{
27    AggregatedStats, AsyncActorLearnerConfig, PPOConfig, PPOTrainerBurn, TrainingStats,
28    compute_entropy_loss, compute_policy_loss, compute_value_loss, generate_minibatch_indices,
29};
30pub use sac::{SacConfig, SacStepStats, SacTrainer};