# How to Use SampleBatch for Trajectory Data Storage and Processing in EvoRL Rollouts

> Learn how to use SampleBatch in EvoRL for efficient trajectory data storage and processing. Optimize rollouts with JAX-compatible PyTree containers for batching and advantage computation.

- Repository: [EMI-Group/evorl](https://github.com/emi-group/evorl)
- Tags: how-to-guide
- Published: 2026-03-01

---

**`SampleBatch` is the core JAX-compatible PyTree container in EvoRL that stores trajectory data from environment interactions, enabling efficient batching, shuffling, and advantage computation across time and batch dimensions.**

EvoRL uses `SampleBatch` as the fundamental data structure for trajectory data storage and processing during rollouts. Defined in [`evorl/sample_batch.py`](https://github.com/emi-group/evorl/blob/main/evorl/sample_batch.py), this container inherits from `PyTreeData` and `PyTreeArrayMixin`, making it fully compatible with JAX transformations like `jax.jit`, `jax.vmap`, and `jax.lax.scan` without requiring manual tree mapping.

## Structure of SampleBatch for Trajectory Storage

The `SampleBatch` class organizes reinforcement learning transitions into a structured PyTree with the following fields:

- **obs**: Observations seen by the agent at the current step, typically shaped `[T, B, ...]` where T represents time steps and B represents batch size.
- **actions**: Actions selected by the policy, matching the observation batch dimensions.
- **rewards**: Scalar rewards received after executing actions, shaped `[T, B]`.
- **next_obs**: Observations at the next time step, used for bootstrapping value estimates.
- **dones**: Episode termination flags where `1` indicates the episode has ended.
- **extras**: Optional dictionary for policy-side or environment-side auxiliary information.

### JAX Compatibility Implementation

In [`evorl/sample_batch.py`](https://github.com/emi-group/evorl/blob/main/evorl/sample_batch.py), the class definition leverages JAX's PyTree protocol:

```python
class SampleBatch(PyTreeData, PyTreeArrayMixin):
    """Data container for trajectory data."""
    obs: chex.ArrayTree | None = None
    actions: chex.ArrayTree | None = None
    rewards: Reward | RewardDict | None = None
    next_obs: chex.Array | None = None
    dones: chex.Array | None = None
    extras: ExtraInfo | None = None

```

This inheritance registers the class as a JAX PyTree node, allowing transformations to traverse its fields automatically without extra boilerplate.

## Creating SampleBatch During EvoRL Rollouts

The trajectory generation process populates `SampleBatch` containers through two primary mechanisms in [`evorl/rollout.py`](https://github.com/emi-group/evorl/blob/main/evorl/rollout.py): single-step transitions and full trajectory collection.

### Single-Step Transition Creation

The `env_step` function constructs individual transition batches representing single environment steps:

```python
transition = SampleBatch(
    obs=env_state.obs,
    actions=actions,
    rewards=env_nstate.reward,
    dones=env_nstate.done,
    next_obs=env_nstate.obs,
    extras=PyTreeDict(
        policy_extras=policy_extras,
        env_extras=env_extras
    ),
)

```

This creates a `SampleBatch` with shape `[B, ...]` where B represents the number of parallel environments.

### Trajectory Aggregation with jax.lax.scan

The `rollout` function aggregates individual transitions into full trajectories using `jax.lax.scan`:

```python
(env_state, _), trajectory = jax.lax.scan(
    _one_step_rollout, 
    (env_state, key), 
    (), 
    length=rollout_length
)

```

The resulting `trajectory` is a `SampleBatch` with shape `[T, B, ...]`, where T is the rollout length. This structure maintains all trajectory data storage and processing in a single, batched container suitable for vectorized operations.

## Post-Processing SampleBatch for Training

EvoRL provides specialized utilities in [`evorl/utils/rl_toolkits.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/rl_toolkits.py) for transforming `SampleBatch` data before training or storage.

### Shuffling and Flattening Trajectories

For replay buffer insertion or epoch-based training, trajectories often require reshaping:

- **`shuffle_sample_batch`**: Randomly permutes the batch dimension of every leaf tensor, useful for breaking temporal correlations in replay buffers.
- **`flatten_rollout_trajectory`**: Collapses time and batch axes from `[T, B, ...]` to `[T·B, ...]`, converting trajectory data into a flat batch suitable for neural network training.
- **`flatten_pop_rollout_episode`**: Flattens population-level batches from `[#pop, T, B, ...]` to `[T, #pop·B, ...]` for evolutionary RL algorithms.

Example usage for flattening a trajectory:

```python
from evorl.utils.rl_toolkits import flatten_rollout_trajectory

# Collapse time and batch dimensions

flat_traj = flatten_rollout_trajectory(trajectory)   # shape (T*B, ...)

```

### Computing Returns and Advantages

The `SampleBatch` structure supports efficient computation of reinforcement learning targets:

```python
from evorl.utils.rl_toolkits import compute_discount_return, compute_gae

# Discounted return per environment

discounted_return = compute_discount_return(
    reward_batch, 
    done_batch, 
    discount=0.99
)

# Generalized Advantage Estimation

value_batch = agent.value_fn(traj.obs)  # Shape (T+1, B)

lambda_returns, advantages = compute_gae(
    rewards=reward_batch,
    values=value_batch,
    dones=done_batch,
    terminations=jnp.zeros_like(done_batch),
    gae_lambda=0.95,
    discount=0.99,
)

```

These operations leverage the PyTree structure to apply transformations consistently across all trajectory fields.

## Multi-Agent SampleBatch Processing

For multi-agent environments processed by [`evorl/rollout_ma.py`](https://github.com/emi-group/evorl/blob/main/evorl/rollout_ma.py), the `SampleBatch` contains joint observations and actions for all agents. Individual agent data extraction uses `tree_get` from [`evorl/utils/tree_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/tree_utils.py):

```python
from evorl.utils.tree_utils import tree_get

# Extract specific agent's data from joint batch

agent_obs = tree_get(sample_batch.obs, agent_id)
agent_actions = tree_get(sample_batch.actions, agent_id)

```

This pattern maintains unified trajectory data storage while enabling independent policy updates for each agent in the population.

## Summary

- **SampleBatch** serves as the fundamental trajectory data storage container in EvoRL, implementing a JAX-compatible PyTree structure defined in [`evorl/sample_batch.py`](https://github.com/emi-group/evorl/blob/main/evorl/sample_batch.py).
- The container stores observations, actions, rewards, next observations, dones, and extras with typical shapes `[T, B, ...]` for time and batch dimensions.
- During rollouts in [`evorl/rollout.py`](https://github.com/emi-group/evorl/blob/main/evorl/rollout.py), `jax.lax.scan` aggregates individual transitions into full trajectories while maintaining the PyTree structure.
- Post-processing utilities in [`evorl/utils/rl_toolkits.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/rl_toolkits.py) provide shuffling, flattening, and advantage estimation functions that operate directly on `SampleBatch` objects.
- Multi-agent rollouts use the same container, extracting per-agent slices via `tree_get` from [`evorl/utils/tree_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/tree_utils.py).

## Frequently Asked Questions

### How does SampleBatch differ between a single transition and a full trajectory?

A **single transition** `SampleBatch` has shape `[B, ...]` containing data from one environment step across B parallel environments. A **full trajectory** `SampleBatch` has shape `[T, B, ...]` where T represents the rollout length. The `rollout` function in [`evorl/rollout.py`](https://github.com/emi-group/evorl/blob/main/evorl/rollout.py) uses `jax.lax.scan` to automatically stack transitions along the time dimension, converting single-step batches into trajectory batches while preserving the PyTree structure.

### Can SampleBatch handle variable-length episodes in a fixed-shape JAX array?

Yes, through the **dones** field and masking utilities. While JAX requires fixed array shapes for JIT compilation, `SampleBatch` uses the `dones` field to mark episode boundaries with `1` values. Functions like `compute_discount_return` and `compute_gae` in [`evorl/utils/rl_toolkits.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/rl_toolkits.py) respect these termination flags to ensure calculations stop at episode boundaries. For storage of variable-length sequences, the `Episode` class in [`evorl/sample_batch.py`](https://github.com/emi-group/evorl/blob/main/evorl/sample_batch.py) provides a `valid_mask` field.

### How do I convert trajectory data for replay buffer storage?

Use the **flattening utilities** in [`evorl/utils/rl_toolkits.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/rl_toolkits.py). First, call `shuffle_sample_batch` to break temporal correlations by randomly permuting the batch dimension. Then use `flatten_rollout_trajectory` to collapse the time and batch dimensions from `[T, B, ...]` to `[T·B, ...]`. This flat structure is compatible with standard replay buffer implementations in [`evorl/replay_buffers/replay_buffer.py`](https://github.com/emi-group/evorl/blob/main/evorl/replay_buffers/replay_buffer.py), which expect uniform batch dimensions for sampling.

### Is SampleBatch compatible with evolutionary RL population-based training?

Yes, through **population-level batching**. EvoRL supports evolutionary algorithms by handling population dimensions in `SampleBatch`. The `flatten_pop_rollout_episode` function in [`evorl/utils/rl_toolkits.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/rl_toolkits.py) reshapes population-level batches from `[#pop, T, B, ...]` to `[T, #pop·B, ...]`, enabling vectorized fitness evaluation across the population while maintaining the trajectory structure required for policy gradient calculations.