Feat: MPE wrapper #1103
Feat: MPE wrapper #1103
Conversation
WiemKhlifi
requested review from
arnupretorius,
DriesSmit,
RuanJohn,
jcformanek,
siddarthsingh1,
sash-a,
OmaymaMahjoub and
ulricharmel
as code owners
| if isinstance(self._env, SimpleSpreadMPE): | ||
| obs, _ = self._env.reset(jax.random.PRNGKey(0)) | ||
| # The shape provided in the `observation_spaces` isn't always the same as | ||
| # that given after constructing the full agent's obs. | ||
| self._env.observation_spaces["agent_0"].shape = obs["agent_0"].shape |
There was a problem hiding this comment.
I think this is fixed in jaxmarl 0.0.5, if that is the case could we please update and remove this 🙏
There was a problem hiding this comment.
Yes this is fixed in that version but we should handle other issues if we fix it to 0.0.5 like respecting a specific version of Jax/ Flax else we'll have a dependencies conflict issue and should handle this error: FLAIROx/JaxMARL#108.
It works for me if jax==0.4.25 and flax should be forced to 0.8 since the one forced by jaxmarl 0.0.5 (0.7.6) gives the same previous issue 😢
There was a problem hiding this comment.
This really sucks. Ye I think they just haven't upgraded flax for their algorithms, which seems to be why they pin to 0.4.25, I'm not really sure what to do here, because I'd really like to pin to the latest jaxmarl and definitely latest jax. @RuanJohn thoughts?
What?
Added a wrapper for MPE environment.
Why?
This change allows to run MPE environments within the Mava ecosystem by supporting both
discreteandcontinuousaction spaces.Extra
Currently supports a fully cooperative
SimpleSpreadMPEtask as the observation and action spaces are homogenous across agents. Further modifications are needed to support other versions.