Mava
Mava copied to clipboard
instadeepai
Feature/jittable agent ids
What?
Change agent ids to use EntityID dataclass - this dataclass can be easily used to represent an agent and is jittable. The dataclass provides utility functions to convert an entityID object to a string and vice versa. The entityID object hashes in such a way that the string representation can be used to index dictionaries and vice versa.
Why?
Strings are not a primitive jax type and as such causes a lot of problems when trying to pass agent information to any jitted function. Additionally, the entityID dataclass provides a standardised way to represent agent ids which allow easy access to information such as type and id. This eliminates the need to perform string manipulation which can be a large source of errors.
How?
An entityId dataclass object was made with the dunder string, dunder hash and dunder equals methods overridden to allow any entityID object to be equally represented as a string. Additionally, utility functions were created to allow users to easily convert between the two. The debugging environment was changed to use entityID objects instead of strings. The mava specs class was changed to use the standardised form of entityID strings.
This PR must go in before PR #496 can be considered. Big change to how Mava handles agent IDs, need approval from @arnupretorius
Just a note because I don't think this will be merged in before I go on leave.
To merge this in we will also have to change env wrappers to be compatible, just two methods need to be changed possible_agents and agents
As an example you would change possible_agents from:
def possible_agents(self):
return [f"agent_{i} for i in range(self.num_agents)
to
def possible_agents(self):
return [f"EntityId(id=i) for i in range(self.num_agents)
I really think this is a nice change to mava, aside from removing one of the bigger obstacles when jitting environments it also adds a lot of stafety when manipulating agent id's as you no longer need to get the substring you can simply call agent_id.id or agent_id.type
I agree
Best Regards, Edan Toledo
From: Sasha @.> Sent: Thursday, July 21, 2022 6:52:16 PM To: instadeepai/Mava @.> Cc: Edan Toledo @.>; Assign @.> Subject: Re: [instadeepai/Mava] Feature/jittable agent ids (PR #491)
CAUTION: This email originated outside the UCT network. Do not click any links or open attachments unless you know and trust the source.
Just a note because I don't think this will be merged in before I go on leave.
To merge this in we will also have to change env wrappers to be compatible, just two methods need to be changed possible_agents and agents
As an example you would change possible_agents from:
def possible_agents(self): return [f"agent_{i} for i in range(self.num_agents)
to
def possible_agents(self): return [f"EntityId(id=i) for i in range(self.num_agents)
I really think this is a nice change to mava, aside from removing one of the bigger obstacles when jitting environments it also adds a lot of stafety when manipulating agent id's as you no longer need to get the substring you can simply call agent_id.id or agent_id.type
— Reply to this email directly, view it on GitHubhttps://eur01.safelinks.protection.outlook.com/?url=https%3A%2F%2Fprotect-za.mimecast.com%2Fs%2FrWXbCAnX51iWKZ1LhGWFwE&data=05%7C01%7Ctldeda001%40myuct.ac.za%7C150b7db816fb4c2dd87708da6b3102d7%7C92454335564e4ccfb0b024445b8c03f7%7C0%7C0%7C637940155525641531%7CUnknown%7CTWFpbGZsb3d8eyJWIjoiMC4wLjAwMDAiLCJQIjoiV2luMzIiLCJBTiI6Ik1haWwiLCJXVCI6Mn0%3D%7C3000%7C%7C%7C&sdata=GhpQD%2By9nmRxQLVm2X16JC9r5Q1IgN0T5NeCdksaPRc%3D&reserved=0, or unsubscribehttps://eur01.safelinks.protection.outlook.com/?url=https%3A%2F%2Fprotect-za.mimecast.com%2Fs%2FJHh-CBgX56f4XDlpf6_RVG&data=05%7C01%7Ctldeda001%40myuct.ac.za%7C150b7db816fb4c2dd87708da6b3102d7%7C92454335564e4ccfb0b024445b8c03f7%7C0%7C0%7C637940155525797285%7CUnknown%7CTWFpbGZsb3d8eyJWIjoiMC4wLjAwMDAiLCJQIjoiV2luMzIiLCJBTiI6Ik1haWwiLCJXVCI6Mn0%3D%7C3000%7C%7C%7C&sdata=Dvgsqbr3jbupsmP2Lx4Idecxnxkp%2Flf7%2FYNR4X56sSc%3D&reserved=0. You are receiving this because you were assigned.Message ID: @.***>
Disclaimer - University of Cape Town This email is subject to UCT policies and email disclaimer published on our website at http://www.uct.ac.za/main/email-disclaimer or obtainable from +27 21 650 9111. If this email is not related to the business of UCT, it is sent by the sender in an individual capacity. Please report security incidents or abuse via https://csirt.uct.ac.za/page/report-an-incident.php.
We are closing this PR as we are focussing on the JAX release. We will re-open it when needed again in the future.
