search

waymo-research / waymax

A JAX-based simulator for autonomous driving research.
Other
831 stars 90 forks source link

The possibility of integrate multi-agent reinforcement learning libraries #3

Open fantastic8124 opened 11 months ago

fantastic8124 commented 11 months ago

Thank you for your excellent work！ From a practical application perspective, I wonder how to integrate the simulator with Marllib or Marl_ Benchmark and other multi-agent reinforcement learning libraries, like other exist simulator like metadrive?

justinjfu commented 11 months ago

Hi @fantastic8124,

It should be possible to write a wrapper to other RL interfaces (we have several examples here: https://github.com/waymo-research/waymax/tree/main/waymax/env/wrappers).

It looks like Marllib uses Ray's MultiAgentEnv interface, which has a sparse representation for agents - actions are passed in as a dict of {agent_key: action} and observations are returned as a dict of {agent_key: obs}. Waymax uses a dense representation by default (we expect and return arrays, with a valid mask), but it should be possible to convert between the two using an adapter.

We'll be busy supporting bugfixes and basic features for the time being with the initial release, but would be open to implementing this in the future or accepting a pull request with this feature.