https://towardsdatascience.com/vectorize-and-parallelize-rl-environments-with-jax-q-learning-at-the-speed-of-light-49d07373adf5/
Vectorize and Parallelize RL Environments with JAX: Q-learning at the Speed of Light⚡ | Towards...
Mar 5, 2025 - Learn to vectorize a GridWorld environment and train 30 Q-learning agents in parallel on a CPU, at 1.8 million step per seconds!
vectorizerlenvironmentsjaxq
https://towardsdatascience.com/breaking-down-state-of-the-art-ppo-implementations-in-jax-6f102c06c149/
A Practical Guide to Proximal Policy Optimization in JAX | Towards Data Science
Mar 5, 2025 - All the tricks and details you wish you knew about PPO
practical guidejax towards
https://towardsdatascience.com/a-gentle-introduction-to-deep-reinforcement-learning-in-jax-c1e45a179b92/
A Gentle Introduction to Deep Reinforcement Learning in JAX | Towards Data Science
Mar 5, 2025 - Solving the CartPole environment with DQN in under a second
gentle introductiondeepjax
https://towardsdatascience.com/implementing-a-transformer-encoder-from-scratch-with-jax-and-haiku-791d31b4f0dd/
Implementing a Transformer Encoder from Scratch with JAX and Haiku 🤖 | Towards Data Science
Mar 5, 2025 - Understanding the fundamental building blocks of Transformers.
implementingtransformerjax