Inquiry About PyTorch Integration and Model Conversion on Waymax

[jxmmy7777]

Hello team,

I'm exploring the integration of PyTorch with Waymax for machine learning tasks. I have a trained model in PyTorch and would like to know the suggested integration approach for using it with the Waymax simulator. Could you provide guidance on integrating PyTorch models with Waymax, especially when the model is already trained in the PyTorch framework?
Any suggestions or references would be appreciated.

Thank you!

[mantralbs]

@jxmmy7777 Hello, have you solve it?

[samuela]

Hi @jxmmy7777! I have converted the issue to a discussion since this is more of a Q&A question. In short: waymax is a JAX library, so you can incorporate it with PyTorch code. A few possibilities include:

1. Using JAX and PyTorch simultaneously. You can convert PyTorch tensors to JAX arrays and vice versa. The easiest way to do this is likely going to be by going through numpy, eg jax_array = jnp.array(torch_array.numpy(force=True)) and torch_tensor = torch.from_numpy(jax_array.asnumpy()). Note that going with this approach means you won't be able to propagate gradients across the JAX/PyTorch boundary without some finagling.
2. Using torch2jax. torch2jax is a conversion layer for running PyTorch code in JAX with native support for all the usual JAX things like vmap, grad, etc.