EdanToledo/Stoix
A research-oriented single-agent reinforcement learning framework implemented entirely in JAX for high-performance distributed training.

Stoix provides optimized implementations of popular single-agent RL algorithms in JAX, enabling easy parallelization across devices via pmap and full compilation with jit for fast training. It targets researchers who want to quickly iterate on RL ideas, offering baseline implementations that can be tuned for specific environments. The codebase supports both JAX-native environments and external environments through its Sebulba system.