IsaacGym JAX integration?

Hi,

I plan to train reinforcement learning with JAX and IsaacGym. However, IsaacGym currently only have PyTorch tensor integration. Is there any method to use JAX tensors in simulation?

Hi there,

I have reached out to the isaac gym team regarding your question.

Cheers,
Andrei

Hi there,

according to the isaac gym team there are currently no plans for the environment implementation in JAX. If it is helpful you could follow the examples of rl-games training using Brax. It is written in JAX, rl-games on PyTorch, and tensor sharing between them happens on the GPU.

Cheers,
Andrei

Thanks!

I will have a look at Brax for RL training examples. By the way, can I use Gym API and manually construct JAX tensors using the device pointers like the gymtorch plugin does?