How does this compare with PyTorch / Tensorflow / etc.? Obviously doing heavy data processing on the GPU will have a large speedup compared to a single thread on the CPU. It's almost like the author is claiming credit for creating Nvidia, when in fact he is just calling its APIs.
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).