thrust_rl/
lib.rs

1//! # Thrust
2//!
3//! High-performance reinforcement learning in Rust.
4//!
5//! Thrust is a modern RL library built on top of the [Burn](https://burn.dev)
6//! tensor framework. After phase 5 of the Burn migration (#65), Burn is the
7//! only tensor backend in the workspace; the previous `tch`/libtorch path
8//! has been removed in favour of Burn's multi-backend stack (CPU NdArray,
9//! WebGPU, CUDA, ROCm, Metal, Vulkan).
10//!
11//! ## Quick Start
12//!
13//! The [`prelude`] re-exports the types you need to stand up a training
14//! run: the [`Environment`](crate::env::Environment) trait and a couple of
15//! built-in environments, the
16//! [`MlpBurnPolicy`](crate::policy::mlp::MlpBurnPolicy) actor-critic network,
17//! and the trainer configs/trainers (A2C, PPO, DQN, SAC, BC). The example below
18//! builds a CartPole env, an A2C config, and a policy on Burn's CPU `NdArray`
19//! backend — the same pieces the `train_cartpole_a2c` example wires into a full
20//! rollout/update loop.
21//!
22//! ```rust
23//! # #[cfg(feature = "training")]
24//! # fn main() {
25//! use burn::backend::{Autodiff, NdArray};
26//! use thrust_rl::prelude::*;
27//!
28//! type Backend = Autodiff<NdArray<f32>>;
29//!
30//! // 1. Build an environment and read its observation / action dims.
31//! let env = CartPole::new();
32//! let obs_dim = env.observation_space().shape[0];
33//! let action_dim = match env.action_space().space_type {
34//!     SpaceType::Discrete(n) => n,
35//!     SpaceType::Box => panic!("CartPole is discrete"),
36//! };
37//!
38//! // 2. Construct the actor-critic policy on the CPU backend.
39//! let device = Default::default();
40//! let policy = MlpBurnPolicy::<Backend>::with_config(
41//!     obs_dim,
42//!     action_dim,
43//!     MlpBurnConfig { hidden_dim: 64, ..Default::default() },
44//!     &device,
45//! );
46//! let _ = policy;
47//!
48//! // 3. Pick a trainer config. See the `train_cartpole_a2c` example for
49//! //    the full rollout-collect + GAE + `A2cTrainer::train_step` loop.
50//! let config = A2cConfig { num_envs: 16, n_steps: 5, ..Default::default() };
51//! assert_eq!(config.num_envs, 16);
52//! # }
53//! # #[cfg(not(feature = "training"))]
54//! # fn main() {}
55//! ```
56
57#![warn(missing_docs)]
58#![warn(clippy::all)]
59
60/// Environment traits and implementations
61pub mod env;
62
63/// Policy and neural network implementations
64/// inference submodule available for WASM, training modules require training
65/// feature
66pub mod policy;
67
68/// Experience buffers and replay management (requires training feature).
69///
70/// The storage layer is plain `Vec<f32>`/`Vec<i64>` regardless of which
71/// tensor backend the trainer ultimately uses; tensor materialization
72/// happens via `to_burn_tensors` on the batch types when the trainer
73/// calls into Burn.
74#[cfg(feature = "training")]
75pub mod buffer;
76
77/// Training algorithms (PPO, DQN).
78#[cfg(feature = "training")]
79pub mod train;
80
81/// Multi-agent training infrastructure (Burn backend).
82///
83/// Synchronized joint trainer plus the multi-agent environment trait and
84/// cross-thread message payloads. Restored on top of the Burn stack in
85/// issue #100 after PR #98 removed the pre-Burn tch-coupled module.
86#[cfg(feature = "training")]
87pub mod multi_agent;
88
89/// Utility functions and helpers
90pub mod utils;
91
92/// Pure Rust inference for WASM compilation
93pub mod inference;
94
95/// WebAssembly bindings for browser visualization
96#[cfg(feature = "wasm")]
97pub mod wasm;
98
99/// Prelude module for convenient imports.
100///
101/// Brings the highest-value public types into scope with a single
102/// `use thrust_rl::prelude::*;`. The set is curated, not a blanket glob:
103///
104/// - **Environment API** (always available): the
105///   [`Environment`](crate::env::Environment) trait plus
106///   [`StepResult`](crate::env::StepResult),
107///   [`SpaceInfo`](crate::env::SpaceInfo),
108///   [`SpaceType`](crate::env::SpaceType), and the built-in
109///   [`CartPole`](crate::env::CartPole) /
110///   [`SimpleBandit`](crate::env::SimpleBandit) environments.
111/// - **Training stack** (requires the `training` feature): the
112///   [`MlpBurnPolicy`](crate::policy::mlp::MlpBurnPolicy) actor-critic network
113///   and its [`MlpBurnConfig`](crate::policy::mlp::MlpBurnConfig), the trainer
114///   configs/trainers (A2C, PPO, DQN, SAC, BC), and the
115///   [`RolloutBuffer`](crate::buffer::rollout::RolloutBuffer) /
116///   [`ReplayBuffer`](crate::buffer::replay::ReplayBuffer) experience buffers.
117///
118/// Training-only re-exports are `#[cfg(feature = "training")]`-gated so
119/// `--no-default-features` still compiles.
120pub mod prelude {
121    // The `crate::env::*` re-export is unconditional; every other line is
122    // `#[cfg(feature = "training")]`-gated so `--no-default-features` still
123    // compiles. (`cargo fmt` alphabetizes the block, so the env import lands
124    // in the middle rather than first.)
125    #[cfg(feature = "training")]
126    pub use crate::buffer::replay::ReplayBuffer;
127    #[cfg(feature = "training")]
128    pub use crate::buffer::rollout::RolloutBuffer;
129    pub use crate::env::{CartPole, Environment, SimpleBandit, SpaceInfo, SpaceType, StepResult};
130    #[cfg(feature = "training")]
131    pub use crate::policy::mlp::{BurnActivation, MlpBurnConfig, MlpBurnPolicy};
132    #[cfg(feature = "training")]
133    pub use crate::train::{
134        A2cConfig, A2cTrainer, BcConfig, BcTrainer, DQNConfig, DQNTrainerBurn, PPOConfig,
135        PPOTrainerBurn, SacConfig, SacTrainer,
136    };
137}
138
139/// Current version of thrust-rl
140pub const VERSION: &str = env!("CARGO_PKG_VERSION");
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145
146    #[test]
147    fn test_version() {
148        assert_eq!(VERSION, "0.3.0");
149    }
150}
thrust_rl/lib.rs

thrust_rl/
lib.rs
