flashbax

Flashbax: High-Speed Buffers in JAX

Overview 🔍

Flashbax is an innovative library designed to enhance the use of experience replay buffers in reinforcement learning (RL). Particularly tailored for compatibility with JAX, Flashbax allows these buffers to be integrated seamlessly into fully compiled functions and training loops. This makes it easy to use different types of buffers such as Flat Buffer, Trajectory Buffer, and Prioritised Buffer variations. Whether you're working on academic research, industrial projects, or personal experiments, Flashbax provides a straightforward and flexible framework for managing RL experience replay.

Features 🛠️

Efficient Buffer Variants: Each buffer type within Flashbax is a special version of a trajectory buffer. This ensures optimal memory use and function across different buffer types.
Flat Buffer: Similar to buffers used in algorithms like DQN, this buffer handles the transition between states efficiently.
Item Buffer: Ideal for storing independent data, like complete episodes or individual tuples of (observation, action, reward, etc.).
Trajectory Buffer: Perfect for sampling multi-step trajectories, particularly helpful for algorithms using recurrent networks.
Prioritised Buffers: These buffers sample based on user-defined priorities, aligning with modern prioritisation techniques.
Trajectory/Flat Queue: Provides a queue for sampling in FIFO order, beneficial for on-policy algorithms.

Setup 🎬

Integrating Flashbax into your projects involves a few simple steps:

Installation: Install Flashbax via pip:
```
pip install flashbax
```

Buffer Selection: Choose the buffer type that fits your needs:

import flashbax as fbx
buffer = fbx.make_trajectory_buffer(...)

Initialize and Use: Initialize the buffer, add data, and sample data easily:

# Initialize
state = buffer.init(example_timestep)
# Add Data
state = buffer.add(state, example_data)
# Sample Data
batch = buffer.sample(state, rng_key)

Quickstart 🏁

Here's a brief example demonstrating how to use the flat buffer:

import jax
import jax.numpy as jnp
import flashbax as fbx

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

# Initialize the buffer state
fake_timestep = {"obs": jnp.array([0, 0]), "reward": jnp.array(0.0)}
state = buffer.init(fake_timestep)

# Add data to the buffer
state = buffer.add(state, {"obs": jnp.array([1, 2]), "reward": jnp.array(3.0)})
state = buffer.add(state, {"obs": jnp.array([4, 5]), "reward": jnp.array(6.0)}) 
state = buffer.add(state, {"obs": jnp.array([7, 8]), "reward": jnp.array(9.0)})

# Sample a transition
rng_key = jax.random.PRNGKey(0)
batch = buffer.sample(state, rng_key)

# Output: obs = [[4 5]], obs' = [[7 8]]
print(f"obs = {batch.experience.first['obs']}, obs' = {batch.experience.second['obs']}")

Examples 🧑‍💻

Flashbax also provides detailed examples in Colab notebooks, covering the setup and usage of various buffer types like Flat Buffers, Trajectory Buffers, and Prioritised Buffers.

Important Considerations ⚠️

When using Flashbax buffers, several important considerations can enhance performance:

Sequential Data: Ensure data is added sequentially for optimal functioning.
Effective Buffer Size: Understand how buffer batch dimensions affect storage.
Episode Truncation: Correctly handle transition between episodes in buffers to prevent data loss.
Independent Data: Consider choosing the appropriate buffer type for independent data use cases.

Benchmarks 📈

Flashbax has been tested and benchmarked against popular buffers, showing competitive performance across CPUs, GPUs, and TPUs, especially in adding and sampling data from buffers.

Contributing 🤝

Flashbax welcomes contributions from the community. Interested contributors can check out the contributing guidelines and the issue tracker for open issues.