instadeepai/flashbax
⚡ Flashbax: Accelerated Replay Buffers in JAX
Flashbax helps reinforcement learning researchers and practitioners manage and sample data efficiently for training AI agents. It takes in sequences of observations, actions, and rewards (experience data) from an agent's interactions and outputs batches of this experience for training algorithms. This is for machine learning engineers and AI researchers building and optimizing reinforcement learning models.
274 stars. No commits in the last 6 months. Available on PyPI.
Use this if you are developing reinforcement learning agents using JAX and need a high-performance, flexible way to store and retrieve past experiences for training.
Not ideal if you are not working with JAX or if your project does not involve reinforcement learning.
Stars
274
Forks
22
Language
Python
License
Apache-2.0
Category
Last pushed
Sep 22, 2025
Commits (30d)
0
Dependencies
7
Get this data via API
curl "https://pt-edge.onrender.com/api/v1/quality/ml-frameworks/instadeepai/flashbax"
Open to everyone — 100 requests/day, no key needed. Get a free key for 1,000/day.
Related frameworks
explosion/thinc
🔮 A refreshing functional take on deep learning, compatible with your favorite libraries
google-deepmind/optax
Optax is a gradient processing and optimization library for JAX.
patrick-kidger/diffrax
Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable....
google/grain
Library for reading and processing ML training data.
patrick-kidger/equinox
Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/