# How to Use the ReplayBuffer for Off-Policy Learning Algorithms in EvoRL

> Learn to use the EvoRL ReplayBuffer for off-policy algorithms. Discover JAX-compatible methods add sample and can_sample for efficient trajectory storage and retrieval for TD3 and SAC.

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

---

**The ReplayBuffer in EvoRL is a functional, JAX-compatible circular buffer that stores environment trajectories and enables uniform random sampling for off-policy algorithms like TD3 and SAC through the `add()`, `sample()`, and `can_sample()` methods.**

The `ReplayBuffer` in the [emi-group/evorl](https://github.com/emi-group/evorl) repository provides a pure-functional storage mechanism designed specifically for off-policy reinforcement learning. Unlike traditional imperative buffers, EvoRL's implementation uses immutable **PyTrees** to maintain state, making it fully compatible with JAX transformations such as `jit`, `grad`, and `pmap` while supporting efficient circular storage for experience replay.

## Core Architecture and Design

EvoRL implements a clean abstraction layer for replay buffer functionality, separating the interface definition from the concrete uniform-sampling implementation.

### AbstractReplayBuffer Interface

The `AbstractReplayBuffer` class in [`evorl/replay_buffers/replay_buffer.py`](https://github.com/emi-group/evorl/blob/main/evorl/replay_buffers/replay_buffer.py) (lines 26-85) defines the required API contract that all buffer implementations must follow. This interface declares five essential methods: `init()` for state creation, `add()` for trajectory insertion, `sample()` for batch retrieval, `can_sample()` for validation, and `is_full()` for capacity checking. By inheriting from this abstract base class, algorithms can swap buffer implementations without modifying their training logic.

### ReplayBufferState PyTree

Buffer state management relies on the `ReplayBufferState` class defined at lines 12-24 of [`evorl/replay_buffers/replay_buffer.py`](https://github.com/emi-group/evorl/blob/main/evorl/replay_buffers/replay_buffer.py). This **PyTree** structure contains three key fields: the data array storing transitions, a `current_index` write pointer, and the `buffer_size` counter tracking filled capacity. The state is purely functional—every operation returns a new state rather than mutating the existing one. For convenient type checking throughout the codebase, `ReplayBufferState` is exported as a type alias (`chex.ArrayTree`) in [`evorl/types.py`](https://github.com/emi-group/evorl/blob/main/evorl/types.py) at line 33.

### Utility Functions

Efficient batched indexing on the stored PyTree is handled by `tree_get` and `tree_set` functions located in [`evorl/utils/jax_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/jax_utils.py). These utilities enable vectorized operations on the nested buffer data without breaking JAX's functional programming constraints.

## Lifecycle of the ReplayBuffer in Off-Policy Training

Off-policy algorithms like TD3 and SAC follow a consistent four-phase lifecycle when interacting with the replay buffer.

### 1. Buffer Initialization

Algorithms instantiate `ReplayBuffer` with three hyperparameters during construction. In [`evorl/algorithms/td3.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/td3.py) (lines 23-29), the buffer is configured as follows:

```python
replay_buffer = ReplayBuffer(
    capacity=config.replay_buffer_capacity,
    min_sample_timesteps=max(config.batch_size, config.learning_start_timesteps),
    sample_batch_size=config.batch_size,
)

```

- **capacity**: Maximum number of timesteps to retain before circular overwrite begins.
- **min_sample_timesteps**: Minimum buffer population required before sampling is permitted (typically set to the larger of batch size or learning start timesteps).
- **sample_batch_size**: Number of transitions returned by each `sample()` call.

The workflow initializes the actual buffer state during setup via `_setup_replaybuffer()` in [`evorl/workflows/rl_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py) (lines 287-292), which calls `buffer.init(sample_spec)` to allocate JAX arrays based on the trajectory structure.

### 2. Adding Trajectories

After each environment rollout, algorithms call the `add()` method to store new experience. As shown in [`evorl/algorithms/td3.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/td3.py) (lines 92-95), the operation follows functional patterns:

```python
replay_buffer_state = self.replay_buffer.add(state.replay_buffer_state, trajectory)

```

The `add()` method writes data into circular storage, updates the `current_index` and `buffer_size` counters, and returns a fresh `ReplayBufferState`. The implementation supports optional binary masks to skip specific elements during insertion, which is essential for handling variable-length trajectories without padding overhead.

### 3. Sampling Batches

When sufficient data is available (`can_sample()` returns `True`), the algorithm draws training batches using a JAX `PRNGKey` for random index generation. In [`evorl/algorithms/td3.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/td3.py) (lines 45-48), the sampling pattern appears as:

```python
sample_batch = self.replay_buffer.sample(replay_buffer_state, rb_key)

```

The `sample()` method generates random indices in the range `[0, buffer_size)` and returns a batch dictionary matching the structure of the original `sample_spec`. Because the buffer state is immutable, the workflow must thread the new state through the training `State` dataclass after each operation.

## Practical Implementation Example

The following example demonstrates creating a buffer, adding dummy trajectories, and sampling training batches:

```python
import chex
import jax
import jax.numpy as jnp
from evorl.replay_buffers import ReplayBuffer

# Define the data structure matching your environment transitions

sample_spec = {
    "obs": jnp.zeros((4,), jnp.float32),
    "action": jnp.zeros((2,), jnp.float32),
    "reward": jnp.zeros((), jnp.float32),
    "discount": jnp.zeros((), jnp.float32),
    "next_obs": jnp.zeros((4,), jnp.float32),
    "done": jnp.zeros((), jnp.bool_),
}

# Initialize buffer with 10k capacity, requiring 1k steps before sampling

buffer = ReplayBuffer(
    capacity=10_000,
    min_sample_timesteps=1_000,
    sample_batch_size=256,
)

state = buffer.init(sample_spec)

# Simulate a rollout with batch dimension 32

rng = jax.random.PRNGKey(0)
dummy_trajectory = {
    "obs": jax.random.normal(rng, (32, 4)),
    "action": jax.random.normal(rng, (32, 2)),
    "reward": jax.random.normal(rng, (32,)),
    "discount": jnp.ones((32,)),
    "next_obs": jax.random.normal(rng, (32, 4)),
    "done": jnp.zeros((32,), jnp.bool_),
}

# Add data and update state functionally

state = buffer.add(state, dummy_trajectory)

# Sample when ready

if buffer.can_sample(state):
    sample_key = jax.random.PRNGKey(42)
    batch = buffer.sample(state, sample_key)
    print({k: v.shape for k, v in batch.items()})  # Each tensor shape: (256, ...)

```

## Integration with EvoRL Algorithms

The `ReplayBuffer` integrates seamlessly with EvoRL's algorithm implementations. Both **TD3** ([[`evorl/algorithms/td3.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/td3.py)](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/td3.py)) and **SAC** ([[`evorl/algorithms/sac.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/sac.py)](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/sac.py)) demonstrate identical buffer interaction patterns:

1. Store trajectories via `self.replay_buffer.add()`
2. Check readiness with `self.replay_buffer.can_sample()`
3. Retrieve batches using `self.replay_buffer.sample()`
4. Feed the resulting `sample_batch` into critic and actor loss functions

This consistency allows researchers to implement new off-policy algorithms by following the established state-threading patterns found in these reference implementations.

## Summary

- **Functional purity** eliminates side effects, enabling JIT compilation (`jit`) and multi-device parallelization (`pmap`) across TPU/GPU clusters.
- **Circular storage** provides constant-time writes and automatic overwriting of oldest data once `capacity` is reached.
- **State threading** is mandatory—always capture the returned `ReplayBufferState` from `add()` operations to maintain data consistency.
- **Mask support** allows selective insertion of trajectory elements, supporting variable-length episodes without expensive copying.
- **Uniform sampling** is implemented in the concrete class, with the `AbstractReplayBuffer` interface permitting custom prioritization strategies.

## Frequently Asked Questions

### How does EvoRL's ReplayBuffer differ from standard replay buffers?

Standard replay buffers typically use imperative, in-place mutations on NumPy arrays. EvoRL's implementation is purely functional, storing data in **immutable PyTrees** that return new state objects on every operation. This design is essential for JAX compatibility, enabling automatic differentiation (`grad`) and compilation (`jit`) while maintaining thread safety across parallel training devices.

### What are the three required parameters when creating a ReplayBuffer?

You must specify `capacity` (maximum stored timesteps), `min_sample_timesteps` (minimum data required before sampling), and `sample_batch_size` (number of transitions per batch). These are typically configured in the algorithm's initialization, as seen in [`evorl/algorithms/td3.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/td3.py) where `min_sample_timesteps` is set to the maximum of batch size and learning start timesteps.

### How do you check if the buffer contains enough data to sample?

Call the `can_sample(buffer_state)` method, which returns a boolean indicating whether the current `buffer_size` exceeds the `min_sample_timesteps` threshold configured during instantiation. This check prevents sampling from partially filled buffers during early training phases.

### Can I implement prioritized experience replay with this buffer?

The concrete `ReplayBuffer` class implements uniform random sampling only. However, the `AbstractReplayBuffer` interface in [`evorl/replay_buffers/replay_buffer.py`](https://github.com/emi-group/evorl/blob/main/evorl/replay_buffers/replay_buffer.py) is designed for extensibility. You can create a custom prioritized buffer by subclassing `AbstractReplayBuffer` and implementing your own `add()` and `sample()` methods that handle importance sampling weights and priority updates.