-
Notifications
You must be signed in to change notification settings - Fork 89
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Feat: MPE wrapper #1103
base: develop
Are you sure you want to change the base?
Feat: MPE wrapper #1103
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looking really good, couple minor things 😄
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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
discrete
andcontinuous
action spaces.Extra
Currently supports a fully cooperative
SimpleSpreadMPE
task as the observation and action spaces are homogenous across agents. Further modifications are needed to support other versions.