Achieving 4000x Speedups with PureJaxRL

chrislu.pagesite

A blog post on how JAX can massively speedup RL training through vectorisation.

from

JAX

added

2026-10-10

likes

0

similar

1. ▲
   purejaxrl github.com

   tool category #jax link ↗

   Vectorisable, end-to-end RL algorithms in JAX.

   
2. ▲
   JAX RL github.com

   tool category #jax link ↗

   Implementations of reinforcement learning algorithms.

   
3. ▲
   Meta-Learning in 50 Lines of JAX by Eric Jang blog.evjang.com

   site category #jax link ↗

   Introduction to both JAX and Meta-Learning.

4. ▲
   Implementing NeRF in JAX by Soumik Rakshit and Saurav Maheshkar wandb.ai

   site category #jax link ↗

   A tutorial on 3D volumetric rendering of scenes represented by Neural Radiance Fields in JAX.

5. ▲
   Pgx github.com

   tool category #jax link ↗

   Vectorized board game environments for RL with an AlphaZero example.

   
6. ▲
   Jax in Action manning.com

   site category #jax link ↗

   A hands-on guide to using JAX for deep learning and other mathematically-intensive applications.

JAX › Tutorials and Blog Posts: “A blog post on how JAX can massively speedup RL training through vectorisation.”