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 sampling::{
23    Minibatch, MinibatchIterator, generate_minibatch_indices, sample_minibatch, shuffle_indices,
24    train_val_split,
25};
26pub use storage::{RolloutBatch, RolloutBuffer, RolloutBurnTensors};
27
28// Submodules
29mod gae;
30mod sampling;
31mod storage;
32
33#[cfg(test)]
34mod tests;
35
36// Legacy interface - re-export compute_advantages as a method on RolloutBuffer
37impl RolloutBuffer {
38    /// Compute advantages using Generalized Advantage Estimation
39    ///
40    /// This is a convenience method that calls the module-level function.
41    /// Iterates the full `[num_steps, num_envs]` capacity; use
42    /// [`Self::compute_advantages_partial`] when the buffer is only
43    /// partially filled.
44    ///
45    /// # Arguments
46    /// * `last_values` - Value estimates for the final states `[num_envs]`
47    /// * `gamma` - Discount factor
48    /// * `gae_lambda` - GAE lambda parameter
49    pub fn compute_advantages(&mut self, last_values: &[f32], gamma: f32, gae_lambda: f32) {
50        gae::compute_advantages(self, last_values, gamma, gae_lambda);
51    }
52
53    /// Compute advantages over the first `valid_steps` rows only.
54    ///
55    /// Convenience wrapper around the module-level
56    /// [`compute_advantages_partial`]. Use this when the rollout buffer
57    /// has been partially filled (`valid_steps < num_steps`) to prevent
58    /// zero-padded tail rows from contaminating GAE on the real prefix.
59    ///
60    /// # Arguments
61    /// * `valid_steps` - Number of filled rows at the start of the buffer
62    /// * `last_values` - Bootstrap `V(s_{T+1})` for the state after row
63    ///   `valid_steps - 1`, per environment `[num_envs]`
64    /// * `gamma` - Discount factor
65    /// * `gae_lambda` - GAE lambda parameter
66    pub fn compute_advantages_partial(
67        &mut self,
68        valid_steps: usize,
69        last_values: &[f32],
70        gamma: f32,
71        gae_lambda: f32,
72    ) {
73        gae::compute_advantages_partial(self, valid_steps, last_values, gamma, gae_lambda);
74    }
75
76    /// Get a batch of all data from the buffer
77    ///
78    /// This is a convenience method that calls the module-level function.
79    /// Returns the full `[num_steps, num_envs]` capacity, including any
80    /// zero-padded unwritten tail; use [`Self::get_filled_batch`] when
81    /// the buffer is only partially filled.
82    pub fn get_batch(&self) -> RolloutBatch {
83        RolloutBatch::from_buffer(self)
84    }
85
86    /// Get a batch covering only the first `valid_steps` rows.
87    ///
88    /// Use this when the buffer was filled with fewer than `num_steps`
89    /// transitions (e.g. an early-terminating rollout). It prevents the
90    /// zero-initialized tail of the buffer from being handed to PPO as
91    /// fake training data.
92    ///
93    /// # Panics
94    /// Panics if `valid_steps > self.shape().0`.
95    pub fn get_filled_batch(&self, valid_steps: usize) -> RolloutBatch {
96        RolloutBatch::from_buffer_partial(self, valid_steps)
97    }
98}
thrust_rl/buffer/rollout.rs

thrust_rl/buffer/
rollout.rs
