-
Notifications
You must be signed in to change notification settings - Fork 5
Expand file tree
/
Copy pathcoma.py
More file actions
executable file
·188 lines (171 loc) · 7.57 KB
/
Copy pathcoma.py
File metadata and controls
executable file
·188 lines (171 loc) · 7.57 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
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
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:
"""
ArgumentsKeys:
- necessary: ``obs`` { ``agent_state``, ``action_mask`` }, ``prev_state``
ReturnsKeys:
- necessary: ``logit``, ``next_state``, ``action_mask``
"""
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
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
"""
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``
"""
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:
COMA network is QAC-type actor-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:
"""
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``
"""
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)