Struct RolloutBuffer 

pub struct RolloutBuffer { /* private fields */ }

Rollout buffer for storing trajectories

Stores trajectories collected from environment interactions and computes advantages using Generalized Advantage Estimation (GAE).

Buffer Layout

The buffer uses a [num_steps, num_envs] layout where:

• num_steps: Number of timesteps per rollout (typically 128-2048)
• num_envs: Number of parallel environments

This layout provides good cache locality for forward passes and efficient computation of advantages.

Implementations

impl RolloutBuffer

pub fn new(num_steps: usize, num_envs: usize, obs_dim: usize) -> Self

Create a new rollout buffer

Arguments

• num_steps - Number of timesteps per rollout
• num_envs - Number of parallel environments
• obs_dim - Dimensionality of observations

pub fn add( &mut self, step: usize, env_id: usize, observation: &[f32], action: i64, reward: f32, value: f32, log_prob: f32, terminated: bool, truncated: bool, )

Add a transition to the buffer

Arguments

• step - Timestep within the rollout (0 to num_steps-1)
• env_id - Environment ID (0 to num_envs-1)
• observation - Current observation
• action - Action taken
• reward - Reward received
• value - Value estimate for current state
• log_prob - Log probability of the action
• terminated - Whether the episode terminated
• truncated - Whether the episode was truncated

pub fn reset(&mut self)

Reset the buffer for a new rollout

pub fn shape(&self) -> (usize, usize, usize)

Get buffer shape (num_steps, num_envs, obs_dim)

pub fn len(&self) -> usize

Get total number of transitions in buffer

pub fn is_empty(&self) -> bool

Check if buffer is empty

pub fn obs_shape(&self) -> (usize, usize)

Get observations tensor shape for neural network input

pub fn observations(&self) -> &[Vec<Vec<f32>>]

Borrow the per-step observations, indexed [step][env] and then by observation dimension (final inner Vec<f32> is one observation vector of length obs_dim).

pub fn actions(&self) -> &[Vec<i64>]

Borrow the per-step discrete actions, indexed [step][env].

pub fn rewards(&self) -> &[Vec<f32>]

Borrow the per-step rewards (raw, un-discounted), indexed [step][env].

pub fn values(&self) -> &[Vec<f32>]

Borrow the per-step bootstrap value estimates V(s_t) produced by the policy at collection time, indexed [step][env]. Used by compute_advantages as the value baseline.

pub fn log_probs(&self) -> &[Vec<f32>]

Borrow the per-step action log-probabilities under the behavior policy at collection time, indexed [step][env]. PPO uses these as the denominator of the importance-sampling ratio.

pub fn terminated(&self) -> &[Vec<bool>]

Borrow the per-step terminal-state flags (true on episode end), indexed [step][env]. GAE zeroes the bootstrap across terminal transitions but keeps the realized reward.

pub fn truncated(&self) -> &[Vec<bool>]

Borrow the per-step truncation flags (true on time-limit / external reset that is not a terminal state), indexed [step][env]. GAE retains the bootstrap value across truncated transitions because the trajectory is still “alive” in the value-function sense.

pub fn advantages(&self) -> &[Vec<f32>]

Borrow the per-step GAE advantages computed by compute_advantages, indexed [step][env]. Zero-initialized until GAE has run.

pub fn returns(&self) -> &[Vec<f32>]

Borrow the per-step value-function targets (advantages + values), indexed [step][env]. Zero-initialized until GAE has run.

pub fn advantages_mut(&mut self) -> &mut [Vec<f32>]

Mutable view of the advantages buffer, indexed [step][env]. Used when in-place normalizing advantages or applying a custom advantage estimator that does not touch returns; use Self::advantages_and_returns_mut when both must be borrowed at the same time (e.g. inside compute_advantages).

pub fn returns_mut(&mut self) -> &mut [Vec<f32>]

Mutable view of the returns / value-target buffer, indexed [step][env]. Use Self::advantages_and_returns_mut instead when both must be borrowed at the same time.

pub fn advantages_and_returns_mut( &mut self, ) -> (&mut [Vec<f32>], &mut [Vec<f32>])

Get mutable references to both advantages and returns This is needed to avoid double mutable borrow in GAE computation

impl RolloutBuffer

pub fn compute_advantages( &mut self, last_values: &[f32], gamma: f32, gae_lambda: f32, )

Compute advantages using Generalized Advantage Estimation

This is a convenience method that calls the module-level function. Iterates the full [num_steps, num_envs] capacity; use Self::compute_advantages_partial when the buffer is only partially filled.

Arguments

• last_values - Value estimates for the final states [num_envs]
• gamma - Discount factor
• gae_lambda - GAE lambda parameter

pub fn compute_advantages_partial( &mut self, valid_steps: usize, last_values: &[f32], gamma: f32, gae_lambda: f32, )

Compute advantages over the first valid_steps rows only.

Convenience wrapper around the module-level compute_advantages_partial. Use this when the rollout buffer has been partially filled (valid_steps < num_steps) to prevent zero-padded tail rows from contaminating GAE on the real prefix.

Arguments

• valid_steps - Number of filled rows at the start of the buffer
• last_values - Bootstrap V(s_{T+1}) for the state after row valid_steps - 1, per environment [num_envs]
• gamma - Discount factor
• gae_lambda - GAE lambda parameter

pub fn compute_vtrace_advantages( &mut self, target_log_probs: &[Vec<f32>], last_values: &[f32], gamma: f32, rho_bar: f32, c_bar: f32, )

Compute V-trace targets and advantages (Espeholt et al. 2018) for off-policy correction.

Convenience wrapper around the module-level compute_vtrace_advantages. The buffer’s stored log_probs are treated as the behavior-policy log-probs; target_log_probs ([num_steps][num_envs]) are the current target policy’s log-probs reevaluated over the stored observations/actions. Results are written into advantages/returns exactly like Self::compute_advantages, so get_batch() works as usual.

Arguments

• target_log_probs - Target-policy log-probs [num_steps][num_envs]
• last_values - Bootstrap V(s_{T+1}) per environment [num_envs]
• gamma - Discount factor
• rho_bar - IS ratio clip for the TD target (typically 1.0)
• c_bar - IS ratio clip for the trace coefficient (typically 1.0)

pub fn get_batch(&self) -> RolloutBatch

Get a batch of all data from the buffer

This is a convenience method that calls the module-level function. Returns the full [num_steps, num_envs] capacity, including any zero-padded unwritten tail; use Self::get_filled_batch when the buffer is only partially filled.

pub fn get_filled_batch(&self, valid_steps: usize) -> RolloutBatch

Get a batch covering only the first valid_steps rows.

Use this when the buffer was filled with fewer than num_steps transitions (e.g. an early-terminating rollout). It prevents the zero-initialized tail of the buffer from being handed to PPO as fake training data.

Panics

Panics if valid_steps > self.shape().0.

Trait Implementations

impl Clone for RolloutBuffer

fn clone(&self) -> RolloutBuffer

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for RolloutBuffer

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.