Async RL in Pure JAX

Async RL in JAX

Scaling RL has become the core approach in unlocking the capabilities of many frontier models today. Thinking beyond current RLVR recipes potentially requires changes across different actors within the stack. Existing frameworks such as Tunix are great for production level RLVR; however, for experimentation, they are quite heavy and require Pathways to scale RL on TPUs. To make this experimentation easier, we introduce a lightweight, performant asynchronous RL stack in pure JAX. We share a work log alongside a few engineering insights we've seen along the way. Code:…

saved by

• Ishaan Panigrahi

related reading

• How to Parallelize a Transformer for Training — an explorable explanationezyang.github.io
• Is Frontier Asynchronous RL Solved? — Luke J. Huangluk-huang.github.io
• Low Latency and Model Training at Modalrhea24.github.io
• State of RL for reasoning LLMs | A. Weersaweers.de
• Keep the Tokens Flowing: Lessons from 16 Open-Source RL Librarieshuggingface.co
• irhum.github.io - Tensor Parallelism with jax.pjitirhum.github.io
• Tiny TPUtinytpu.com
• [2510.13786] The Art of Scaling Reinforcement Learning Compute for LLMsarxiv.org
• [2602.17616] Stable Asynchrony: Variance-Controlled Off-Policy RL for LLMsarxiv.org
• The 4-bitter Lesson | humans&humansand.ai
• Rethinking RL Infra for Agents | B'Logbillxbf.github.io
• IsoCompute Playbook: Optimally Scaling Sampling Compute for RL Training of LLMscompute-optimal-rl-llm-scaling.github.io