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?