Where do you actually jit the functions? parallel/batched hyperparameter search similar to PureJaxRL?

[ugurbolat]

Hi! Really cool that you took the ideas from PureJaxRL to next level.

I searched the keyword "jit" in the codebase and nothing shows up 😅 . Are the training loops jitted similar to PureJaxRL or left to the user?

I was considering implementing a hyperparameter search similar to PureJaxRL where training loops are jitted and vmapped. Is the current state too early for this?

[EdanToledo]

Hey, thanks. I don't explicitly use jax.jit since I use pmap to distribute over multiple (or just one if you only have one) devices. Pmapping in JAX will automatically JIT your functions for you.

Regarding the hyperparameter search, no it's not too early to do something like this in fact for some research of mine I did this over multiple seeds however I will say that since I use config dicts, you would need to slightly hack the input of the train function to take hyperparameters that you can vmap over. Alternatively, everything in Stoix is ready to use with Optuna Bayesian optimisation so you can more intelligently tune parameters. The benefit of this is that although it's not done in parallel, it efficiently explores the space of possible hyperparameters rather than you having to define the specific values that you care about. I personally find this more efficient than choosing my own set of values especially for algorithms I'm not as familiar with and even more when I'm tuning multiple hyperparameters at once since a multidimensional vmap over a large grid is not the most efficient.

[ugurbolat]

Completely missed that pmap will jit the function. thanks for the pointer.

you would need to slightly hack the input of the train function to take hyperparameters that you can vmap over.

I took a close look and I am not able to see this as a slight hack

In PureJaxRL implementation, train function pushes almost everything inside such as network param. initialization, optimizer, env. initialization so that it is vmapped over rng/key/seed.

One possible point of entry for the hack is to learn function which is pmapped originally. However, all the initialization and instantiation process above for network param., optimizer, env. etc are now outside of the scope of learn function. Therefore, for learn function to be batched, we need to take care of what is left outside. For example,

• we would to need to add additional dimension to network and to env.
• vmap/vectorize _update_step function as well.

Or split learner_setup function into somehow that would like to how train function defined in PureJaxRL implementation.

Does that make sense? Or am I completely off with my current thinking?

What got me existed about PureJaxRL is the ability to parallelize/batch the whole training loop (i.e., the whole agent) in a single GPU. This is because I would like to do a research on the hyperparameter seach methods and what optuna offers would be comparison to other methods in the end.

[EdanToledo]

Hello, so you're not wrong and the point of entry would be the learn_fn but its not that hard to do. Essentially, since we already have the update_batch_axis which is specified in the anakin arch config, we can repurpose that for hyperparameters. All you would need to do is add the hyperparameter you want to vmap over to the learner state and then when broadcasting the learner state to be of size (num_devices, update_batch_size, ...) you simply put in the hyperparameters values you care about in there in the correct shape instead of replicating it for all values. Subsequently, after that all you need to do is remove the pmean of grads (since each run is now independent) and then slightly tweak the logger to log all runs independently. I've done this for another project as ive said and it worked very nicely. The hacking is maybe 20-30 lines of code and it mainly (75%+) has to do with the logging actually more than the difficulty in changing the learn fn.

[EdanToledo]

If this is something that people would be interested in, i dont mind posting a small notebook tutorial of how one would do it but im not sure i have the time currently.