Thread by @pcastr

📢📢Dopamine now also runs on #JAX !!📢📢

Happy to announce that our #RL library, dopamine, now also has JAX implementations of all our agents, including a new agent: QR-DQN!
github.com/google/dopamine
1/X

The JAX philosophy goes very well with that of Dopamine: flexibility for research without sacrificing simplicity.
I've been using it for a while now and I've found its modularity quite appealing in terms of simplifying some of the more difficult aspects of the agents.
2/X

Consider the projection operator used for the C51 and Rainbow agents, which is a rather complex TF op. This went from 23 lines of fairly complex TF code to 9 lines of JAX code that is more straightforward.
3/X

We are able to achieve this simplicity in large part because of JAX's vmap functionality, which helps us avoid explicitly dealing with batch dimensions.
jax.readthedocs.io/en/latest/jax.…
4/X

Thanks to JAX's Just in Time compilation our JAX agents achieve the same speed of training as our TensorFlow agents.
jax.readthedocs.io/en/latest/jax.…
5/X

In addition to the existing agents, we also provide a new agent: QR-DQN!
This has been a common request for some time now.
github.com/google/dopamin…
6/X

We provide trained checkpoints for all our JAX agents:
github.com/google/dopamin…
7/X

We also provide a new colab notebook where you can generate videos of your trained agents!
colab.research.google.com/github/google/…
8/X

We look forward to seeing all the new and exciting results this extra code provides. Happy JAXing! 🙂
9/9