My Authors
Read all threads
📢📢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
Missing some Tweet in this thread? You can try to force a refresh.

Keep Current with Pablo Samuel Castro

Profile picture

Stay in touch and get notified when new unrolls are available from this author!

Read all threads

This Thread may be Removed Anytime!

Twitter may remove this content at anytime, convert it as a PDF, save and print for later use!

Try unrolling a thread yourself!

how to unroll video

1) Follow Thread Reader App on Twitter so you can easily mention us!

2) Go to a Twitter thread (series of Tweets by the same owner) and mention us with a keyword "unroll" @threadreaderapp unroll

You can practice here first or read more on our help page!

Follow Us on Twitter!

Did Thread Reader help you today?

Support us! We are indie developers!


This site is made by just two indie developers on a laptop doing marketing, support and development! Read more about the story.

Become a Premium Member ($3.00/month or $30.00/year) and get exclusive features!

Become Premium

Too expensive? Make a small donation by buying us coffee ($5) or help with server cost ($10)

Donate via Paypal Become our Patreon

Thank you for your support!