Tutorials
Tutorials for learning JAX and Reinforcement Learning.
A Gentle Introduction to Deep Reinforcement Learning in JAX
Solving the CartPole environment with DQN in under a second
Achieving Over 4000x Speedups and Meta-Evolving Discoveries with PureJaxRL
We can leverage recent advancements in JAX to train parallelised RL agents over 4000x faster entirely on GPUs. Unlike past RL implementations, ours is written end-to-end in Jax.
Introduction to JAX by Google
A beginner-friendly introduction to JAX from the official team.
Spinning Up in Deep RL
Spinning Up is an educational resource produced by OpenAI that makes it easier to learn about deep reinforcement learning (deep RL).
Vectorize and Parallelize RL Environments with JAX: Q-learning at the Speed of Light
Learn to vectorize a GridWorld environment and train 30 Q-learning agents in parallel on a CPU, at 1.8 million step per seconds!
Writing an RL Environment in JAX
How to run CartPole at 1.25 Billion Step/Sec