Flashbax

Vault is an efficient mechanism for saving Flashbax buffers to persistent storage, specifically designed for offline reinforcement learning.

It handles large datasets by extending the temporal/sequential dimension ($T$) of a buffer to a virtually unconstrained degree. It works by reading and writing slices of buffers along the temporal axis, allowing gigantic buffer stores to reside on disk while sub-buffers are loaded into RAM/VRAM for efficient training. Vault supports the item, flat, and trajectory buffers.

Flashbax is a JAX-compatible library for experience replay buffers in reinforcement learning. It provides several buffer variants designed to work within fully compiled functions (jax.jit) and training loops (jax.pmap).

Core buffer types include:

• Flat Buffer: Similar to transition buffers used in DQN; stores sequences of $(s_t, s_{t+1})$ with a period of 1.
• Item Buffer: Stores individual items (e.g., observation, action, reward tuples) independent of each other.
• Trajectory Buffer: Facilitates sampling multi-step trajectories, useful for recurrent networks (e.g., R2D2).
• Prioritised Buffers: Prioritised versions of both Flat and Trajectory buffers, allowing sampling based on user-defined priorities (aligned with PER principles).
• Trajectory/Flat Queue: A FIFO queue structure for on-policy algorithms.

Flashbax buffers use a block-like structure where the effective size depends on the add_batch_size. When using the max_size argument (which sets the total desired number of timesteps), the internal time axis length is calculated as max_size // add_batch_size.

Because this uses integer division (rounding down), your effective buffer size might be smaller than expected. To ensure accuracy, you should either:

1. Use the max_length (time axis) argument explicitly.
2. Ensure the max_size argument is a multiple of the add_batch_size.

If you are evaluating different replay buffer implementations, the following libraries are frequently compared or used in the community:

• Dejax: JAX-compatible replay buffers.
• Reverb: Efficient replay buffers for local and large-scale distributed RL.
• Dopamine: Research framework for fast prototyping with core replay buffers.
• StableBaselines3: RL baselines suite with easy-to-use replay buffers.

Vault extends the temporal axis of a Flashbax buffer state to store experience data to disk. While Vault handles bookkeeping (only saving new, unseen data), Flashbax states are implemented as ring buffers.

Critical Requirement: You must update the Vault frequently enough to prevent unseen data from being overwritten in the ring buffer. If your buffer state has a time-axis length of $\tau$, you must save to the Vault at least every $\tau - 1$ steps.

To integrate Flashbax into your project, install it using pip:

pip install flashbax

Flashbax buffers follow a functional pattern. You first create a buffer object using a make_* function, then use its pure functions to manage state. The buffer object itself is a container for these functions.

Workflow:

1. Create: Use fbx.make_flat_buffer(...), fbx.make_trajectory_buffer(...), etc.
2. Initialise: Call buffer.init(example_data) to create the initial state.
3. Add Data: Call state = buffer.add(state, new_data) to update the state.
4. Sample: Call batch = buffer.sample(state, rng_key) to retrieve data.

import flashbax as fbx

# 1. Create
buffer = fbx.make_flat_buffer(max_length=32, min_length=2, sample_batch_size=1)

# 2. Initialise
state = buffer.init(example_timestep)

# 3. Add Data
state = buffer.add(state, example_data)

# 4. Sample
batch = buffer.sample(state, rng_key)

Flashbax provides high-speed buffers implemented in JAX, designed for efficient data storage and sampling in reinforcement learning or large-scale simulations.

To use Flashbax, you typically follow a three-step workflow:

1. Initialize a buffer with a specific capacity and data structure.
2. Add Data to the buffer using the buffer's API.
3. Sample Data from the buffer for training or evaluation.

Flashbax is built to be highly performant by leveraging JAX's capabilities.

# Typical workflow pattern:
# 1. Initialize
# 2. Add Data
# 3. Sample Data

Because Flashbax buffers are often large and occupy significant device memory, you should perform in-place updates to avoid the overhead of creating copies. To enable this, you must specify the buffer state argument for donation when using jax.jit via the donate_argnums parameter.

If you omit donate_argnums, JAX will force a copy of the state for any modifications, which can negate performance benefits.

def train(train_state, buffer_state):
    ...
    return train_state, buffer_state

# Initialise the buffer state
buffer_fn = fbx.make_trajectory_buffer(...)
buffer_state = buffer_fn.init(example_timestep)

# Initialise some training state
train_state = train_state.init(...)

# Compile the training function and specify the donation of the buffer state argument
train_state, buffer_state = jax.jit(train, donate_argnums=(1,))(
    train_state, buffer_state
)

To use the example notebooks and dependencies provided in the repository, install the package with the [examples] extra. You may also need to clone the repository first if you are working locally with the source code.

!git clone https://github.com/instadeepai/flashbax.git
!pip install ./flashbax[examples]