[Question] Best way to implement stochastic reinforcement learning actors? about flax HOT 2 CLOSED

Hi!

Basically something like (1) is the canonical way we do it! We use a with nn.stochastic(prng_key): context around the model evaluation inside the training-step function, with a top-level prng-key fed into the training-step function from an outside split in the training loop. The important thing is that the training-step function is jitted in its entirety - the one thing not to do is to jit across a nn.stochastic context. Provided it's used inside a jit and the prng-keys are fed like any other function argument, there should be no trouble using it. Using the stochastic context can save a lot of prng-key plumbing boilerplate in models.

I've typed up a quick demo of the canonical way in a colab at https://colab.research.google.com/drive/1eDXEVd8NPXgaSwn7jEMxsZTDVHNUtHYK

Let me know if that helps or if anything remains unclear!

from flax.

Thank you for the amazing answer, I will get working on my project right away! Everything is clear, you even mentioned the part about passing the PRNG key itself while initializing a Module following 3, which something I had problems with.

I'm closing the issue, thanks for the help!

from flax.