FLAIROx / JaxMARL

Multi-Agent Reinforcement Learning with JAX
Apache License 2.0
393 stars 68 forks source link

Add example of animation with MAPPO trained agent #64

Closed satpreetsingh closed 6 months ago

satpreetsingh commented 6 months ago

Would it be possible to add a simple example animation of a MAPPO trained agent?

I don't see it here: https://colab.research.google.com/github/FLAIROx/JaxMARL/blob/main/jaxmarl/tutorials/JaxMARL_Walkthrough.ipynb

Nor do I see it in any of the baselines etc.

Ideally, the agent will be saved to disk and reloaded for generating the animation. Thank you!

amacrutherford commented 6 months ago

Hey! Providing weights for pre-trained agents is on our TODO list but not a priority. You can generate a gif of the MAPPO agent using the training script provided in the baselines folder

satpreetsingh commented 6 months ago

Sorry I meant to ask for a simple example of inference using a MAPPO trained agent. The animation/save/reload were frills that were easy to figure out from your existing code.

I'd like to see your recommended (or canonical) way to reuse the agent/network classes outside the training loop. Thanks!

amacrutherford commented 5 months ago

ah sure, we use safetensors, can example can be see in the IQL script (same code should work for MAPPO). An example of these parameter's being loaded and used is given here. Hope this helps!