Closed alexunderch closed 9 months ago
Hello,
No, it would be very challenging to make the environments stateless and implemented exclusively in JAX. That said, there's no reason why you couldn't use JAX for the agent side, and that's what we do internally. We use something similar to the A3C setup where there are processes stepping the environment in CPU only, and the agent is doing inference and learning on the accelerator (GPU/TPU).
By the looks of things, it the only feasible variant right now. Thank you for the response. Closing the issue.
Are there any ways to natively wrap an environment to be suitable for running a rl algorithm with jax (make the environment stateless), or is it planned to add such support?
Thank you!