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

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 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 (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. 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 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. 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 (lines 23-29), the buffer is configured as follows:

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 (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 (lines 92-95), the operation follows functional patterns:

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 (lines 45-48), the sampling pattern appears as:

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:

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)) and SAC ([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 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 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}