[Feature] Introduce grouping in VMAS #1658
Conversation
facebook-github-bot
added
the
CLA Signed
| @@ -50,9 +63,75 @@ def _get_envs(): | |||
| ] | |||
|
|
|||
|
|
|||
| @set_gym_backend("gym") | |||
There was a problem hiding this comment.
Detached vmas from gym spec conversion to have custom conversion rules
| agent_indices = {} | ||
| action_list = [] | ||
| n_agents = 0 | ||
| for group, agent_names in self.group_map.items(): | ||
| group_action = tensordict.get((group, "action")) | ||
| group_action_list = list(self.read_action(group_action, group=group)) | ||
| agent_indices.update( | ||
| { | ||
| self.agent_names_to_indices_map[agent_name]: i + n_agents | ||
| for i, agent_name in enumerate(agent_names) | ||
| } | ||
| ) | ||
| n_agents += len(agent_names) | ||
| action_list += group_action_list | ||
| action = [action_list[agent_indices[i]] for i in range(self.n_agents)] |
There was a problem hiding this comment.
this will convert actions from group dicts to the list (ordered as it should)
This bit of logic is the focus of the new test
| group_reward_spec, | ||
| group_info_spec, | ||
| ) = self._make_unbatched_group_specs(group) | ||
| self.unbatched_action_spec[group] = group_action_spec |
There was a problem hiding this comment.
is self.unbatched_action_spec public on purpose?
Can we review the public attributes of VMAS? I feel there are a lot of them and most are undocumented. This will hamper development in the future. I would make anything that is not user facing private to give us more dev freedom.
There was a problem hiding this comment.
Sure I will make private all I don't need to be public.
This will be bc-breaking though as the methods will not be available anymore.
Regarding the unbatched_spec, these are needed when creating modules. Giving the specs with the batch_size to the modules will lead to errors #1018. These attirbutes are used in BenchMARL for example
There was a problem hiding this comment.
This will be bc-breaking though as the methods will not be available anymore.
you can always make a property that raises a warning
torchrl/envs/libs/vmas.py
Outdated
| @@ -419,14 +590,10 @@ def read_reward(self, rewards): | |||
| rewards = _selective_unsqueeze(rewards, batch_size=self.batch_size) | |||
| return rewards | |||
|
|
|||
| def read_action(self, action): | |||
| def read_action(self, action, group: str): | |||
There was a problem hiding this comment.
this is a public method and we're adding a new arg without default value, hence it's bc-breaking. Any chance we can make this non-bc breaking?
There was a problem hiding this comment.
I can, or I can make it private, as requested above, which will be more bc breaking though
Co-authored-by: Vincent Moens <vincentmoens@gmail.com>
|
Ok I made all methods private. I would like to keep the attributes public so that they can be accessed by libraries like BenchMARL. |
vmoens
changed the title
torchrl/envs/libs/vmas.py
Outdated
| rewards = _selective_unsqueeze(rewards, batch_size=self.batch_size) | ||
| return rewards | ||
|
|
||
| def read_action(self, action): | ||
| def _read_action(self, action, group: str): |
There was a problem hiding this comment.
I think this solution is worse, it used to be public and now it's private. Is there a solution that does not break previous behaviour?
There was a problem hiding this comment.
yes but it will leave the method public.
Just so to understand, do you want these methods public or private? I am happy to leave them public and making thewm bc comptible, but I thoght that private is what you wanted
There was a problem hiding this comment.
The least bc-breaking the better.
I wasn't pointing at these methods in particular, more at all the attributes created when creating specs. It seems like there are loads of them.
read_action is public in other wrappers so I don't think it's a problem if it's public here too.
|
ok now there shold be no bc changes and the public attributes should be all doxumented |
This PR allows VMAS to use the MARL api with grouping.
Users will be able to provide a group_map to VMAS.
By default this will try to group agents named
name_ninto a group named "name".In case at least one agent does not follow this convention and not group map is specified, VMAS will try to group all agents under the "agent" group