Re:4000x Speedup in Reinforcement Learning with Jax
by luchris429 on Saturday April 08, 2023 @11:42AM
(#35494376)
The baseline we are comparing to is standard RL training that is widely used in academia. The technique mentioned in the blog post is not widely used amongst researchers.
The reason we write about Jax is that doing this technique is really hard in PyTorch / Tensorflow. This is because:
1. Jax has vmap. (PyTorch does now too, but it is far more recent).
2. There are RL environments that others have written in pure Jax (see the blog post for four different repos of RL environments)
3. As m00x hints to, Jax replicates Numpy's API. This makes it way easier to use for non-neural network programming (e.g. RL environments).
The reason we write about Jax is that doing this technique is really hard in PyTorch / Tensorflow. This is because:
1. Jax has vmap. (PyTorch does now too, but it is far more recent).
2. There are RL environments that others have written in pure Jax (see the blog post for four different repos of RL environments)
3. As m00x hints to, Jax replicates Numpy's API. This makes it way easier to use for non-neural network programming (e.g. RL environments).