from typing import Dict, Union import torch import torch.nn as nn from functools import reduce from ding.torch_utils import one_hot, MLP from ding.utils import squeeze, list_split, MODEL_REGISTRY, SequenceType from .q_learning import DRQN class COMAActorNetwork(nn.Module): """ Overview: Decentralized actor network in COMA algorithm. Interface: ``__init__``, ``forward`` """ def __init__( self, obs_shape: int, action_shape: int, hidden_size_list: SequenceType = [128, 128, 64], ): """ Overview: Initialize COMA actor network Arguments: - obs_shape (:obj:`int`): the dimension of each agent's observation state - action_shape (:obj:`int`): the dimension of action shape - hidden_size_list (:obj:`list`): the list of hidden size, default to [128, 128, 64] """ super(COMAActorNetwork, self).__init__() self.main = DRQN(obs_shape, action_shape, hidden_size_list) def forward(self, inputs: Dict) -> Dict: """ Overview: The forward computation graph of COMA actor network Arguments: - inputs (:obj:`dict`): input data dict with keys ['obs', 'prev_state'] - agent_state (:obj:`torch.Tensor`): each agent local state(obs) - action_mask (:obj:`torch.Tensor`): the masked action - prev_state (:obj:`torch.Tensor`): the previous hidden state Returns: - output (:obj:`dict`): output data dict with keys ['logit', 'next_state', 'action_mask'] ArgumentsKeys: - necessary: ``obs`` { ``agent_state``, ``action_mask`` }, ``prev_state`` ReturnsKeys: - necessary: ``logit``, ``next_state``, ``action_mask`` Examples: >>> T, B, A, N = 4, 8, 3, 32 >>> embedding_dim = 64 >>> action_dim = 6 >>> data = torch.randn(T, B, A, N) >>> model = COMAActorNetwork((N, ), action_dim, [128, embedding_dim]) >>> prev_state = [[None for _ in range(A)] for _ in range(B)] >>> for t in range(T): >>> inputs = {'obs': {'agent_state': data[t], 'action_mask': None}, 'prev_state': prev_state} >>> outputs = model(inputs) >>> logit, prev_state = outputs['logit'], outputs['next_state'] """ agent_state = inputs['obs']['agent_state'] prev_state = inputs['prev_state'] if len(agent_state.shape) == 3: # B, A, N agent_state = agent_state.unsqueeze(0) unsqueeze_flag = True else: unsqueeze_flag = False T, B, A = agent_state.shape[:3] agent_state = agent_state.reshape(T, -1, *agent_state.shape[3:]) prev_state = reduce(lambda x, y: x + y, prev_state) output = self.main({'obs': agent_state, 'prev_state': prev_state, 'enable_fast_timestep': True}) logit, next_state = output['logit'], output['next_state'] next_state, _ = list_split(next_state, step=A) logit = logit.reshape(T, B, A, -1) if unsqueeze_flag: logit = logit.squeeze(0) return {'logit': logit, 'next_state': next_state, 'action_mask': inputs['obs']['action_mask']} class COMACriticNetwork(nn.Module): """ Overview: Centralized critic network in COMA algorithm. Interface: ``__init__``, ``forward`` """ def __init__( self, input_size: int, action_shape: int, hidden_size: int = 128, ): """ Overview: initialize COMA critic network Arguments: - input_size (:obj:`int`): the size of input global observation - action_shape (:obj:`int`): the dimension of action shape - hidden_size_list (:obj:`list`): the list of hidden size, default to 128 Returns: - output (:obj:`dict`): output data dict with keys ['q_value'] Shapes: - obs (:obj:`dict`): ``agent_state``: :math:`(T, B, A, N, D)`, ``action_mask``: :math:`(T, B, A, N, A)` - prev_state (:obj:`list`): :math:`[[[h, c] for _ in range(A)] for _ in range(B)]` - logit (:obj:`torch.Tensor`): :math:`(T, B, A, N, A)` - next_state (:obj:`list`): :math:`[[[h, c] for _ in range(A)] for _ in range(B)]` - action_mask (:obj:`torch.Tensor`): :math:`(T, B, A, N, A)` """ super(COMACriticNetwork, self).__init__() self.action_shape = action_shape self.act = nn.ReLU() self.mlp = nn.Sequential( MLP(input_size, hidden_size, hidden_size, 2, activation=self.act), nn.Linear(hidden_size, action_shape) ) def forward(self, data: Dict) -> Dict: """ Overview: forward computation graph of qmix network Arguments: - data (:obj:`dict`): input data dict with keys ['obs', 'prev_state', 'action'] - agent_state (:obj:`torch.Tensor`): each agent local state(obs) - global_state (:obj:`torch.Tensor`): global state(obs) - action (:obj:`torch.Tensor`): the masked action ArgumentsKeys: - necessary: ``obs`` { ``agent_state``, ``global_state`` }, ``action``, ``prev_state`` ReturnsKeys: - necessary: ``q_value`` Examples: >>> agent_num, bs, T = 4, 3, 8 >>> obs_dim, global_obs_dim, action_dim = 32, 32 * 4, 9 >>> coma_model = COMACriticNetwork( >>> obs_dim - action_dim + global_obs_dim + 2 * action_dim * agent_num, action_dim) >>> data = { >>> 'obs': { >>> 'agent_state': torch.randn(T, bs, agent_num, obs_dim), >>> 'global_state': torch.randn(T, bs, global_obs_dim), >>> }, >>> 'action': torch.randint(0, action_dim, size=(T, bs, agent_num)), >>> } >>> output = coma_model(data) """ x = self._preprocess_data(data) q = self.mlp(x) return {'q_value': q} def _preprocess_data(self, data: Dict) -> torch.Tensor: """ Overview: preprocess data to make it can be used by MLP net Arguments: - data (:obj:`dict`): input data dict with keys ['obs', 'prev_state', 'action'] - agent_state (:obj:`torch.Tensor`): each agent local state(obs) - global_state (:obj:`torch.Tensor`): global state(obs) - action (:obj:`torch.Tensor`): the masked action ArgumentsKeys: - necessary: ``obs`` { ``agent_state``, ``global_state``} , ``action``, ``prev_state`` Return: - x (:obj:`torch.Tensor`): the data can be used by MLP net, including \ ``global_state``, ``agent_state``, ``last_action``, ``action``, ``agent_id`` """ t_size, batch_size, agent_num = data['obs']['agent_state'].shape[:3] agent_state_ori, global_state = data['obs']['agent_state'], data['obs']['global_state'] # splite obs, last_action and agent_id agent_state = agent_state_ori[..., :-self.action_shape - agent_num] last_action = agent_state_ori[..., -self.action_shape - agent_num:-agent_num] last_action = last_action.reshape(t_size, batch_size, 1, -1).repeat(1, 1, agent_num, 1) agent_id = agent_state_ori[..., -agent_num:] action = one_hot(data['action'], self.action_shape) # T, B, A，N action = action.reshape(t_size, batch_size, -1, agent_num * self.action_shape).repeat(1, 1, agent_num, 1) action_mask = (1 - torch.eye(agent_num).to(action.device)) action_mask = action_mask.view(-1, 1).repeat(1, self.action_shape).view(agent_num, -1) # A, A*N action = (action_mask.unsqueeze(0).unsqueeze(0)) * action # T, B, A, A*N global_state = global_state.unsqueeze(2).repeat(1, 1, agent_num, 1) x = torch.cat([global_state, agent_state, last_action, action, agent_id], -1) return x @MODEL_REGISTRY.register('coma') class COMA(nn.Module): """ Overview: The network of COMA algorithm, which is QAC-type actor-critic. Interface: ``__init__``, ``forward`` Properties: - mode (:obj:`list`): The list of forward mode, including ``compute_actor`` and ``compute_critic`` """ mode = ['compute_actor', 'compute_critic'] def __init__( self, agent_num: int, obs_shape: Dict, action_shape: Union[int, SequenceType], actor_hidden_size_list: SequenceType ) -> None: """ Overview: initialize COMA network Arguments: - agent_num (:obj:`int`): the number of agent - obs_shape (:obj:`Dict`): the observation information, including agent_state and \ global_state - action_shape (:obj:`Union[int, SequenceType]`): the dimension of action shape - actor_hidden_size_list (:obj:`SequenceType`): the list of hidden size """ super(COMA, self).__init__() action_shape = squeeze(action_shape) actor_input_size = squeeze(obs_shape['agent_state']) critic_input_size = squeeze(obs_shape['agent_state']) + squeeze(obs_shape['global_state']) + \ agent_num * action_shape + (agent_num - 1) * action_shape critic_hidden_size = actor_hidden_size_list[-1] self.actor = COMAActorNetwork(actor_input_size, action_shape, actor_hidden_size_list) self.critic = COMACriticNetwork(critic_input_size, action_shape, critic_hidden_size) def forward(self, inputs: Dict, mode: str) -> Dict: """ Overview: forward computation graph of COMA network Arguments: - inputs (:obj:`dict`): input data dict with keys ['obs', 'prev_state', 'action'] - agent_state (:obj:`torch.Tensor`): each agent local state(obs) - global_state (:obj:`torch.Tensor`): global state(obs) - action (:obj:`torch.Tensor`): the masked action ArgumentsKeys: - necessary: ``obs`` { ``agent_state``, ``global_state``, ``action_mask`` }, ``action``, ``prev_state`` ReturnsKeys: - necessary: - compute_critic: ``q_value`` - compute_actor: ``logit``, ``next_state``, ``action_mask`` Shapes: - obs (:obj:`dict`): ``agent_state``: :math:`(T, B, A, N, D)`, ``action_mask``: :math:`(T, B, A, N, A)` - prev_state (:obj:`list`): :math:`[[[h, c] for _ in range(A)] for _ in range(B)]` - logit (:obj:`torch.Tensor`): :math:`(T, B, A, N, A)` - next_state (:obj:`list`): :math:`[[[h, c] for _ in range(A)] for _ in range(B)]` - action_mask (:obj:`torch.Tensor`): :math:`(T, B, A, N, A)` - q_value (:obj:`torch.Tensor`): :math:`(T, B, A, N, A)` Examples: >>> agent_num, bs, T = 4, 3, 8 >>> agent_num, bs, T = 4, 3, 8 >>> obs_dim, global_obs_dim, action_dim = 32, 32 * 4, 9 >>> coma_model = COMA( >>> agent_num=agent_num, >>> obs_shape=dict(agent_state=(obs_dim, ), global_state=(global_obs_dim, )), >>> action_shape=action_dim, >>> actor_hidden_size_list=[128, 64], >>> ) >>> prev_state = [[None for _ in range(agent_num)] for _ in range(bs)] >>> data = { >>> 'obs': { >>> 'agent_state': torch.randn(T, bs, agent_num, obs_dim), >>> 'action_mask': None, >>> }, >>> 'prev_state': prev_state, >>> } >>> output = coma_model(data, mode='compute_actor') >>> data= { >>> 'obs': { >>> 'agent_state': torch.randn(T, bs, agent_num, obs_dim), >>> 'global_state': torch.randn(T, bs, global_obs_dim), >>> }, >>> 'action': torch.randint(0, action_dim, size=(T, bs, agent_num)), >>> } >>> output = coma_model(data, mode='compute_critic') """ assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode) if mode == 'compute_actor': return self.actor(inputs) elif mode == 'compute_critic': return self.critic(inputs)