Top AI Repos — open-source AI, indexed and scored
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
🏛️A research-friendly codebase for fast experimentation of single-agent reinforcement learning in JAX • End-to-End JAX RL
| Date | Stars |
|---|---|
| 2026-07-31 | 416 |
| 2026-08-02 | 416 |
| 2026-08-03 | 417 |
| 2026-08-06 | 417 |
Today
— stars today
This week
— stars this week
This month
— stars this month
Momentum
0.0
growth rate 0.00%/day
<p align="center">
<a href="docs/images/stoix.png">
<img src="docs/images/stoix.png" alt="Stoix logo" width="30%"/>
</a>
</p>
<div align="center">
<a href="https://www.python.org/doc/versions/">
<img src="https://img.shields.io/badge/python-3.10-blue" alt="Python Versions">
</a>
<a href="https://github.com/instadeepai/Mava/blob/main/LICENSE">
<img src="https://img.shields.io/badge/License-Apache%202.0-orange.svg" alt="License" />
</a>
<a href="https://github.com/psf/black">
<img src="https://img.shields.io/badge/code%20style-black-000000.svg" alt="Code Style" />
</a>
<a href="http://mypy-lang.org/">
<img src="https://www.mypy-lang.org/static/mypy_badge.svg" alt="MyPy" />
</a>
<a href="https://zenodo.org/doi/10.5281/zenodo.10916257"><img src="https://zenodo.org/badge/758685996.svg" alt="DOI"></a>
</div>
<h2 align="center">
<p>Distributed Single-Agent Reinforcement Learning End-to-End in JAX</p>
</h2>
<div align="center">
**_stoic - a person who can endure pain or hardship without showing their feelings or complaining._**
</div>
## Welcome to Stoix! 🏛️
Stoix provides simplified code for quickly iterating on ideas in single-agent reinforcement learning with useful implementations of popular single-agent RL algorithms in JAX allowing for easy parallelisation across devices with JAX's `pmap`. All implementations are fully compiled with JAX's `jit` thus making training and environment execution very fast. However, this does require environments written in JAX. For environments not written in JAX, Stoix offers Sebulba systems (see below). Algorithms and their default hyperparameters have not been optimised for any specific environment and are useful as a starting point for research and/or for initial baselines before hyperparameter tuning.
To join us in these efforts, please feel free to reach out, raise issues or read our [contribution guidelines](#contributing-) (or just star 🌟 to stay up to date with the latest developments)!
Stoix is fully in JAX with substantial speed improvement compared to other popular libraries. We currently provide native support for the [Jumanji][jumanji] environment API and wrappers for popular RL environments.
## System Design Paradigms
Stoix offers two primary system design paradigms (Podracer Architectures) to cater to different research and deployment needs:
- **Anakin:** Traditional Stoix implementations are fully end-to-end compiled with JAX, focusing on speed and simplicity with native JAX environments. This design paradigm is ideal for setups where all components, including environments, can be optimized using JAX, leveraging the full power of JAX's pmap and jit. For an illustration of the Anakin architecture, see this [figure](docs/images/anakin_arch.jpg) from the [Mava][mava] technical report or the original podracer [paper][anakin_paper].
- **Sebulba:** The Sebulba system introduces flexibility by allowing different devices to be assigned specifically for learning and acting. In this setup, acting devices serve as inference servers for multiple parallel environments, which can be written in any framework, not just JAX. Each set of parallel environments can run on separate threads asynchronously. This enables Stoix to be used with a broader range of environments while still benefiting from JAX's speed. For an illustration of the Sebulba architecture, see this [animation](docs/images/sebulba_arch.gif) from the [InstaDeep Sebulba implementation](https://github.com/instadeepai/sebulba/). Whilst one could optimisie for throughput of data to purely maximise performance, Stoix aims to serve as a research codebase, thus, we take inspiration from [cleanba][cleanba] which focused on ensuring reproducibility and algorithm correctness. With this in mind, one could easily adapt the sebulba systems to optimise for throughput.
Not all implementations have both Anakin and Sebulba implementations but effort has gone into making the two implementations as sExcerpt of 19,253 characters
Read on GitHubWould you bet a product on this? Bounded 0–100 and slow moving.
matched fp:0847da1a0a1b51dc, topic:reinforcement-learning, topic:deep-reinforcement-learning, desc:reinforcement learning