thrust_rl/buffer/
replay.rs

1//! Replay buffer for off-policy training (DQN).
2//!
3//! This module implements two flavors of experience replay used by
4//! [`crate::train::dqn`]:
5//!
6//! 1. [`ReplayBuffer`] — a fixed-capacity FIFO buffer with uniform sampling.
7//!    The classic DQN recipe.
8//! 2. [`PrioritizedReplayBuffer`] — a sum-tree backed buffer with proportional
9//!    priority sampling and importance-sampling correction weights (Schaul et
10//!    al. 2015).
11//! 3. [`ContinuousReplayBuffer`] — the continuous-action sibling of
12//!    [`ReplayBuffer`], storing `Vec<f32>` actions of width `action_dim` for
13//!    off-policy continuous control (SAC / TD3 / DDPG). Same ring/FIFO design
14//!    and flat-`Vec` storage; sample it via [`sample_continuous`].
15//!
16//! All store transitions as flat CPU-side `Vec`s rather than Burn tensors
17//! so the buffer stays WASM-compatible if it ever needs to be exposed
18//! there. Tensor materialization happens via `to_burn_tensors` on the
19//! batch types when the trainer pushes them to a device.
20//!
21//! # Quick example (uniform)
22//!
23//! ```ignore
24//! use thrust_rl::buffer::replay::{ReplayBuffer, sample};
25//! use rand::SeedableRng;
26//!
27//! let mut buf = ReplayBuffer::new(/* capacity */ 50_000, /* obs_dim */ 4);
28//! buf.push(&[0.0; 4], /* action */ 1, /* reward */ 1.0, &[0.1; 4], /* done */ false);
29//!
30//! let mut rng = rand::rngs::StdRng::seed_from_u64(0);
31//! if buf.is_ready(/* min_size */ 1) {
32//!     let batch = sample(&buf, 64, &mut rng);
33//!     let device = Default::default();
34//!     let t = batch.to_burn_tensors::<MyBackend>(&device);
35//!     // ... feed t.* into the DQN trainer ...
36//! }
37//! ```
38//!
39//! # Quick example (prioritized)
40//!
41//! ```ignore
42//! use thrust_rl::buffer::replay::PrioritizedReplayBuffer;
43//! use rand::SeedableRng;
44//!
45//! let mut buf = PrioritizedReplayBuffer::new(50_000, 4, /* α */ 0.6, /* ε */ 1e-6);
46//! buf.push(&[0.0; 4], 1, 1.0, &[0.1; 4], false);
47//!
48//! let mut rng = rand::rngs::StdRng::seed_from_u64(0);
49//! if buf.is_ready(1) {
50//!     let batch = buf.sample(64, /* β */ 0.4, &mut rng);
51//!     let device = Default::default();
52//!     let t = batch.to_burn_tensors::<MyBackend>(&device);
53//!     // ... compute weighted Smooth-L1 loss, backprop ...
54//!     // ... then write the new TD-error magnitudes back:
55//!     // buf.update_priorities(&batch.indices, &new_td_errors);
56//! }
57//! ```
58
59pub use continuous_sampling::{
60    ContinuousReplayBatch, ContinuousReplayBurnTensors, sample as sample_continuous,
61};
62pub use continuous_storage::ContinuousReplayBuffer;
63pub use prioritized::{PrioritizedBatch, PrioritizedBurnTensors, PrioritizedReplayBuffer};
64pub use sampling::{ReplayBatch, ReplayBurnTensors, sample};
65pub use storage::ReplayBuffer;
66pub use sum_tree::SumTree;
67
68mod continuous_sampling;
69mod continuous_storage;
70mod prioritized;
71mod sampling;
72mod storage;
73mod sum_tree;
thrust_rl/buffer/replay.rs

thrust_rl/buffer/
replay.rs
