GPU-Based Environment

[DavidAkinpelu]

I am trying to create a GPU-based environment where a model is being trained say resnet18 where the number of environment can be greater than 1. I am not familiar with Jax but I am planning to learn it. However, I want to see how feasible this is with PureJaxrl before I start this voyage.

[DavidAkinpelu]

@luchris429 Any advice?

[luchris429]

That sounds reasonable, but I really can't tell without more details.

[DavidAkinpelu]

@luchris429 Thanks for your response. I want to create an adaptive learning rate controller agent. I want the the environment step to involve training and evaluating a model. The action at the end of each step is to adjust the learning rate. How can I easily parallelize this for multiple environments?