Example: DQN about flax HOT 8 CLOSED

Hi @joaogui1! It would be great to see a DQN example, and perhaps it's even a good idea to use RLAX (if RLAX is general purpose enough to work with any JAX codebase)

But I'm not sure we'll be taking this into flax/examples/. We want to encourage the flourishing of an ecosystem of examples, and not all of them can or should be held to the very high maintenance expectations of flax/examples.

We'll start linking to many GitHub repositories of Flax projects from our README, so my suggestion is to go ahead and build your DQN repo and we'll link to it. Then we could potentially consider bringing that into flax/examples but we'd need one of the Flax OWNERS to take responsibility for it, which involves performance monitoring, commitment to modifications with future API changes, and more.

from flax.

Actually, in particular I would be interested to see what a DQN example with or without RLAX looks like, is that something that would be pretty straightforward to try out?

from flax.

So @avital , I already have the version without rlax (the "tense" version if you will): colab
Though I need some help making it prettier/more idiomatic, specially with the act part which I ended up commenting
I will code a version with rlax and try to put here before the weekend

from flax.

Here's the rlax version @avital

from flax.

Awesome @joaogui1 ! Your examples looks very cool and it would be great if we could link to this from our examples page. Would you be willing to put your example in a separate Github repository, with some instructions on what it is and how to run it? Then we will add a link to it from our example page.

from flax.

Like this @marcvanzee ?

from flax.

Just pinging @marcvanzee

from flax.

@marcvanzee -- can you add this to the table in examples/README.md? https://github.com/joaogui1/RL-JAX/tree/master/DQN

from flax.