thrust_rl/buffer/
rollout.rs

1//! Rollout buffer for storing and processing trajectories
2//!
3//! This module implements experience storage for PPO training, including:
4//! - Trajectory storage (observations, actions, rewards, etc.)
5//! - Generalized Advantage Estimation (GAE) computation
6//! - Efficient sampling for minibatch training
7//!
8//! # Buffer Layout
9//!
10//! The buffer uses a `[num_steps, num_envs]` layout where:
11//! - `num_steps`: Number of timesteps per rollout (typically 128-2048)
12//! - `num_envs`: Number of parallel environments
13//!
14//! This layout provides good cache locality for forward passes and
15//! efficient computation of advantages.
16
17// Re-export main components
18pub use gae::{
19    compute_advantages, compute_advantages_multi_agent, compute_advantages_partial,
20    compute_mc_returns, compute_nstep_returns, normalize_advantages,
21};
22pub use recurrent::{RecurrentMinibatchIterator, RecurrentRolloutBatch, RecurrentRolloutBuffer};
23pub use sampling::{
24    Minibatch, MinibatchIterator, generate_minibatch_indices, sample_minibatch, shuffle_indices,
25    train_val_split,
26};
27pub use storage::{RolloutBatch, RolloutBuffer, RolloutBurnTensors};
28pub use vtrace::compute_vtrace_advantages;
29
30// Submodules
31mod gae;
32mod recurrent;
33mod sampling;
34mod storage;
35mod vtrace;
36
37#[cfg(test)]
38mod tests;
39
40// Legacy interface - re-export compute_advantages as a method on RolloutBuffer
41impl RolloutBuffer {
42    /// Compute advantages using Generalized Advantage Estimation
43    ///
44    /// This is a convenience method that calls the module-level function.
45    /// Iterates the full `[num_steps, num_envs]` capacity; use
46    /// [`Self::compute_advantages_partial`] when the buffer is only
47    /// partially filled.
48    ///
49    /// # Arguments
50    /// * `last_values` - Value estimates for the final states `[num_envs]`
51    /// * `gamma` - Discount factor
52    /// * `gae_lambda` - GAE lambda parameter
53    pub fn compute_advantages(&mut self, last_values: &[f32], gamma: f32, gae_lambda: f32) {
54        gae::compute_advantages(self, last_values, gamma, gae_lambda);
55    }
56
57    /// Compute advantages over the first `valid_steps` rows only.
58    ///
59    /// Convenience wrapper around the module-level
60    /// [`compute_advantages_partial`]. Use this when the rollout buffer
61    /// has been partially filled (`valid_steps < num_steps`) to prevent
62    /// zero-padded tail rows from contaminating GAE on the real prefix.
63    ///
64    /// # Arguments
65    /// * `valid_steps` - Number of filled rows at the start of the buffer
66    /// * `last_values` - Bootstrap `V(s_{T+1})` for the state after row
67    ///   `valid_steps - 1`, per environment `[num_envs]`
68    /// * `gamma` - Discount factor
69    /// * `gae_lambda` - GAE lambda parameter
70    pub fn compute_advantages_partial(
71        &mut self,
72        valid_steps: usize,
73        last_values: &[f32],
74        gamma: f32,
75        gae_lambda: f32,
76    ) {
77        gae::compute_advantages_partial(self, valid_steps, last_values, gamma, gae_lambda);
78    }
79
80    /// Compute V-trace targets and advantages (Espeholt et al. 2018) for
81    /// off-policy correction.
82    ///
83    /// Convenience wrapper around the module-level
84    /// [`compute_vtrace_advantages`]. The buffer's stored `log_probs` are
85    /// treated as the behavior-policy log-probs; `target_log_probs`
86    /// (`[num_steps][num_envs]`) are the current target policy's log-probs
87    /// reevaluated over the stored observations/actions. Results are
88    /// written into `advantages`/`returns` exactly like
89    /// [`Self::compute_advantages`], so `get_batch()` works as usual.
90    ///
91    /// # Arguments
92    /// * `target_log_probs` - Target-policy log-probs `[num_steps][num_envs]`
93    /// * `last_values` - Bootstrap `V(s_{T+1})` per environment `[num_envs]`
94    /// * `gamma` - Discount factor
95    /// * `rho_bar` - IS ratio clip for the TD target (typically 1.0)
96    /// * `c_bar` - IS ratio clip for the trace coefficient (typically 1.0)
97    pub fn compute_vtrace_advantages(
98        &mut self,
99        target_log_probs: &[Vec<f32>],
100        last_values: &[f32],
101        gamma: f32,
102        rho_bar: f32,
103        c_bar: f32,
104    ) {
105        vtrace::compute_vtrace_advantages(
106            self,
107            target_log_probs,
108            last_values,
109            gamma,
110            rho_bar,
111            c_bar,
112        );
113    }
114
115    /// Get a batch of all data from the buffer
116    ///
117    /// This is a convenience method that calls the module-level function.
118    /// Returns the full `[num_steps, num_envs]` capacity, including any
119    /// zero-padded unwritten tail; use [`Self::get_filled_batch`] when
120    /// the buffer is only partially filled.
121    pub fn get_batch(&self) -> RolloutBatch {
122        RolloutBatch::from_buffer(self)
123    }
124
125    /// Get a batch covering only the first `valid_steps` rows.
126    ///
127    /// Use this when the buffer was filled with fewer than `num_steps`
128    /// transitions (e.g. an early-terminating rollout). It prevents the
129    /// zero-initialized tail of the buffer from being handed to PPO as
130    /// fake training data.
131    ///
132    /// # Panics
133    /// Panics if `valid_steps > self.shape().0`.
134    pub fn get_filled_batch(&self, valid_steps: usize) -> RolloutBatch {
135        RolloutBatch::from_buffer_partial(self, valid_steps)
136    }
137}
thrust_rl/buffer/rollout.rs

thrust_rl/buffer/
rollout.rs
