luchris429
While I understand the frustration, I assure you that the rest of the code is functional. This was a simple oversight and should be a trivial fix. I appreciate your feedback and understanding.
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).
While there is work on putting RL environments on accelerators, the main speedup from this work comes from also training many RL agents in parallel. This is largely because the neural networks we use in RL are relatively small and thus don't utilize the GPU very efficiently.
While this was always possible to do, Jax makes it way easier because we just need to call `jax.vmap` to get it to work.
https://github.com/Unity-Technologies/ml-agents
It makes it really easy to make games for reinforcement learning. I worked on it a bit over summer, so if you have any questions, feel free to reach out to me.