0

I am trying to train a DQN agent in an environment coded in JAX, but the initialization of the trainer fails when it first tries to reset the environment (with a not-valid JAX type error). Before getting into the debugging process, I thought of looking for example projects but I cannot find anything so i am wondering whether interfacing them is not possible.

blindeyes
  • 409
  • 3
  • 13

1 Answers1

0

Currently there is no example in RLlib right out of the box that supports jax envs. I think the main functionalities should be there, it was just never tested with Jax envs to figure out the feature requirements and the missing parts. I encourage you to give it a try and create a feature request for it on ray's issue tracker on github and if you have bandwidth give it a try and make a contribution. You can also let the RL team know about the missing features for enabling RLlib to work on JAX envs.

Kourosh Hakhamaneshi
  • 11
  • As it’s currently written, your answer is unclear. Please [edit] to add additional details that will help others understand how this addresses the question asked. You can find more information on how to write good answers [in the help center](/help/how-to-answer). – Community Oct 08 '22 at 04:31