← all repositories

EdanToledo/Stoix

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

410 stars Python ML Frameworks
Stoix
Velocity · 7d
+0.5
★ / day
Trend
steady
star history

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.

heatdrop uses Google Analytics to see which pages get read — nothing else. Your call. How we handle data.