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