Source code for ding.model.template.collaq
from typing import Union, List
import torch
import torch.nn as nn
import torch.nn.functional as F
from functools import reduce
from ding.utils import list_split, MODEL_REGISTRY
from ding.torch_utils import fc_block, MLP, ScaledDotProductAttention
from .q_learning import DRQN
from .qmix import Mixer
class CollaQMultiHeadAttention(nn.Module):
"""
Overview:
The head of collaq attention module.
Interface:
``__init__``, ``forward``
"""
def __init__(
self,
n_head: int,
d_model_q: int,
d_model_v: int,
d_k: int,
d_v: int,
d_out: int,
dropout: float = 0.,
activation: nn.Module = nn.ReLU()
):
"""
Overview:
initialize the head of collaq attention module
Arguments:
- n_head (:obj:`int`): the num of head
- d_model_q (:obj:`int`): the size of input q
- d_model_v (:obj:`int`): the size of input v
- d_k (:obj:`int`): the size of k, used by Scaled Dot Product Attention
- d_v (:obj:`int`): the size of v, used by Scaled Dot Product Attention
- d_out (:obj:`int`): the size of output q
- dropout (:obj:`float`): Dropout ratio, defaults to 0.
- activation (:obj:`nn.Module`): Activation in FFN after attention.
"""
super(CollaQMultiHeadAttention, self).__init__()
self.act = activation
self.n_head = n_head
self.d_k = d_k
self.d_v = d_v
self.w_qs = nn.Linear(d_model_q, n_head * d_k)
self.w_ks = nn.Linear(d_model_v, n_head * d_k)
self.w_vs = nn.Linear(d_model_v, n_head * d_v)
self.fc1 = fc_block(n_head * d_v, n_head * d_v, activation=self.act)
self.fc2 = fc_block(n_head * d_v, d_out)
self.attention = ScaledDotProductAttention(d_k=d_k)
self.layer_norm_q = nn.LayerNorm(n_head * d_k, eps=1e-6)
self.layer_norm_k = nn.LayerNorm(n_head * d_k, eps=1e-6)
self.layer_norm_v = nn.LayerNorm(n_head * d_v, eps=1e-6)
def forward(self, q, k, v, mask=None):
"""
Overview:
forward computation graph of collaQ multi head attention net.
Arguments:
- q (:obj:`torch.nn.Sequential`): the transformer information q
- k (:obj:`torch.nn.Sequential`): the transformer information k
- v (:obj:`torch.nn.Sequential`): the transformer information v
Returns:
- q (:obj:`torch.nn.Sequential`): the transformer output q
- residual (:obj:`torch.nn.Sequential`): the transformer output residual
Shapes:
- q (:obj:`torch.nn.Sequential`): :math:`(B, L, N)` where B is batch_size, L is sequence length, \
N is the size of input q
- k (:obj:`torch.nn.Sequential`): :math:`(B, L, N)` where B is batch_size, L is sequence length, \
N is the size of input k
- v (:obj:`torch.nn.Sequential`): :math:`(B, L, N)` where B is batch_size, L is sequence length, \
N is the size of input v
- q (:obj:`torch.nn.Sequential`): :math:`(B, L, N)` where B is batch_size, L is sequence length, \
N is the size of output q
- residual (:obj:`torch.nn.Sequential`): :math:`(B, L, N)` where B is batch_size, L is sequence length, \
N is the size of output residual
Examples:
>>> net = CollaQMultiHeadAttention(1, 2, 3, 4, 5, 6)
>>> q = torch.randn(1, 2, 2)
>>> k = torch.randn(1, 3, 3)
>>> v = torch.randn(1, 3, 3)
>>> q, residual = net(q, k, v)
"""
d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
batch_size, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)
# Pass through the pre-attention projection: batch_size x len_q x (n_head * d_v)
# Separate different heads: batch_size x len_q x n_head x d_v
q = self.w_qs(q).view(batch_size, len_q, n_head, d_k)
k = self.w_ks(k).view(batch_size, len_k, n_head, d_k)
v = self.w_vs(v).view(batch_size, len_v, n_head, d_v)
residual = q
# Transpose for attention dot product: batch_size x n_head x len_q x d_v
q, k, v = self.layer_norm_q(q).transpose(1, 2), self.layer_norm_k(k).transpose(
1, 2
), self.layer_norm_v(v).transpose(1, 2)
# Unsqueeze the mask tensor for head axis broadcasting
if mask is not None:
mask = mask.unsqueeze(1)
q = self.attention(q, k, v, mask=mask)
# Transpose to move the head dimension back: batch_size x len_q x n_head x d_v
# Combine the last two dimensions to concatenate all the heads together: batch_size x len_q x (n*dv)
q = q.transpose(1, 2).contiguous().view(batch_size, len_q, -1)
q = self.fc2(self.fc1(q))
return q, residual
class CollaQSMACAttentionModule(nn.Module):
"""
Overview:
Collaq attention module. Used to get agent's attention observation. It includes agent's observation\
and agent's part of the observation information of the agent's concerned allies
Interface:
``__init__``, ``_cut_obs``, ``forward``
"""
def __init__(
self,
q_dim: int,
v_dim: int,
self_feature_range: List[int],
ally_feature_range: List[int],
attention_size: int,
activation: nn.Module = nn.ReLU()
):
"""
Overview:
initialize collaq attention module
Arguments:
- q_dim (:obj:`int`): the dimension of transformer output q
- v_dim (:obj:`int`): the dimension of transformer output v
- self_features (:obj:`torch.Tensor`): output self agent's attention observation
- ally_features (:obj:`torch.Tensor`): output ally agent's attention observation
- attention_size (:obj:`int`): the size of attention net layer
- activation (:obj:`nn.Module`): Activation in FFN after attention.
"""
super(CollaQSMACAttentionModule, self).__init__()
self.self_feature_range = self_feature_range
self.ally_feature_range = ally_feature_range
self.attention_layer = CollaQMultiHeadAttention(
1, q_dim, v_dim, attention_size, attention_size, attention_size, activation=activation
)
def _cut_obs(self, obs: torch.Tensor):
"""
Overview:
cut the observed information into self's observation and allay's observation
Arguments:
- obs (:obj:`torch.Tensor`): input each agent's observation
Returns:
- self_features (:obj:`torch.Tensor`): output self agent's attention observation
- ally_features (:obj:`torch.Tensor`): output ally agent's attention observation
Shapes:
- obs (:obj:`torch.Tensor`): :math:`(T, B, A, N)` where T is timestep, B is batch_size, \
A is agent_num, N is obs_shape
- self_features (:obj:`torch.Tensor`): :math:`(T, B, A, N)` where T is timestep, B is batch_size, \
A is agent_num, N is self_feature_range[1] - self_feature_range[0]
- ally_features (:obj:`torch.Tensor`): :math:`(T, B, A, N)` where T is timestep, B is batch_size, \
A is agent_num, N is ally_feature_range[1] - ally_feature_range[0]
"""
# obs shape = (T, B, A, obs_shape)
self_features = obs[:, :, :, self.self_feature_range[0]:self.self_feature_range[1]]
ally_features = obs[:, :, :, self.ally_feature_range[0]:self.ally_feature_range[1]]
return self_features, ally_features
def forward(self, inputs: torch.Tensor):
"""
Overview:
forward computation to get agent's attention observation information
Arguments:
- obs (:obj:`torch.Tensor`): input each agent's observation
Returns:
- obs (:obj:`torch.Tensor`): output agent's attention observation
Shapes:
- obs (:obj:`torch.Tensor`): :math:`(T, B, A, N)` where T is timestep, B is batch_size, \
A is agent_num, N is obs_shape
"""
# obs shape = (T, B ,A, obs_shape)
obs = inputs
self_features, ally_features = self._cut_obs(obs)
T, B, A, _ = self_features.shape
self_features = self_features.reshape(T * B * A, 1, -1)
ally_features = ally_features.reshape(T * B * A, A - 1, -1)
self_features, ally_features = self.attention_layer(self_features, ally_features, ally_features)
self_features = self_features.reshape(T, B, A, -1)
ally_features = ally_features.reshape(T, B, A, -1)
# note: we assume self_feature is near the ally_feature here so we can do this concat
obs = torch.cat(
[
obs[:, :, :, :self.self_feature_range[0]], self_features, ally_features,
obs[:, :, :, self.ally_feature_range[1]:]
],
dim=-1
)
return obs
[docs]@MODEL_REGISTRY.register('collaq')
class CollaQ(nn.Module):
"""
Overview:
The network of CollaQ (Collaborative Q-learning) algorithm.
It includes two parts: q_network and q_alone_network.
The q_network is used to get the q_value of the agent's observation and \
the agent's part of the observation information of the agent's concerned allies.
The q_alone_network is used to get the q_value of the agent's observation and \
the agent's observation information without the agent's concerned allies.
Multi-Agent Collaboration via Reward Attribution Decomposition
https://arxiv.org/abs/2010.08531
Interface:
``__init__``, ``forward``, ``_setup_global_encoder``
"""
def __init__(
self,
agent_num: int,
obs_shape: int,
alone_obs_shape: int,
global_obs_shape: int,
action_shape: int,
hidden_size_list: list,
attention: bool = False,
self_feature_range: Union[List[int], None] = None,
ally_feature_range: Union[List[int], None] = None,
attention_size: int = 32,
mixer: bool = True,
lstm_type: str = 'gru',
activation: nn.Module = nn.ReLU(),
dueling: bool = False,
) -> None:
"""
Overview:
Initialize Collaq network.
Arguments:
- agent_num (:obj:`int`): the number of agent
- obs_shape (:obj:`int`): the dimension of each agent's observation state
- alone_obs_shape (:obj:`int`): the dimension of each agent's observation state without\
other agents
- global_obs_shape (:obj:`int`): the dimension of global observation state
- action_shape (:obj:`int`): the dimension of action shape
- hidden_size_list (:obj:`list`): the list of hidden size
- attention (:obj:`bool`): use attention module or not, default to False
- self_feature_range (:obj:`Union[List[int], None]`): the agent's feature range
- ally_feature_range (:obj:`Union[List[int], None]`): the agent ally's feature range
- attention_size (:obj:`int`): the size of attention net layer
- mixer (:obj:`bool`): use mixer net or not, default to True
- lstm_type (:obj:`str`): use lstm or gru, default to gru
- activation (:obj:`nn.Module`): Activation function in network, defaults to nn.ReLU().
- dueling (:obj:`bool`): use dueling head or not, default to False.
"""
super(CollaQ, self).__init__()
self.attention = attention
self.attention_size = attention_size
self._act = activation
self.mixer = mixer
if not self.attention:
self._q_network = DRQN(
obs_shape, action_shape, hidden_size_list, lstm_type=lstm_type, dueling=dueling, activation=activation
)
else:
# TODO set the attention layer here beautifully
self._self_attention = CollaQSMACAttentionModule(
self_feature_range[1] - self_feature_range[0],
(ally_feature_range[1] - ally_feature_range[0]) // (agent_num - 1),
self_feature_range,
ally_feature_range,
attention_size,
activation=activation
)
# TODO get the obs_dim_after_attention here beautifully
obs_shape_after_attention = self._self_attention(
# torch.randn(
# 1, 1, (ally_feature_range[1] - ally_feature_range[0]) //
# ((self_feature_range[1] - self_feature_range[0])*2) + 1, obs_dim
# )
torch.randn(1, 1, agent_num, obs_shape)
).shape[-1]
self._q_network = DRQN(
obs_shape_after_attention,
action_shape,
hidden_size_list,
lstm_type=lstm_type,
dueling=dueling,
activation=activation
)
self._q_alone_network = DRQN(
alone_obs_shape,
action_shape,
hidden_size_list,
lstm_type=lstm_type,
dueling=dueling,
activation=activation
)
embedding_size = hidden_size_list[-1]
if self.mixer:
self._mixer = Mixer(agent_num, global_obs_shape, embedding_size, activation=activation)
self._global_state_encoder = nn.Identity()
[docs] def forward(self, data: dict, single_step: bool = True) -> dict:
"""
Overview:
The forward method calculates the q_value of each agent and the total q_value of all agents.
The q_value of each agent is calculated by the q_network, and the total q_value is calculated by the mixer.
Arguments:
- data (:obj:`dict`): input data dict with keys ['obs', 'prev_state', 'action']
- agent_state (:obj:`torch.Tensor`): each agent local state(obs)
- agent_alone_state (:obj:`torch.Tensor`): each agent's local state alone, \
in smac setting is without ally feature(obs_along)
- global_state (:obj:`torch.Tensor`): global state(obs)
- prev_state (:obj:`list`): previous rnn state, should include 3 parts: \
one hidden state of q_network, and two hidden state if q_alone_network for obs and obs_alone inputs
- action (:obj:`torch.Tensor` or None): if action is None, use argmax q_value index as action to\
calculate ``agent_q_act``
- single_step (:obj:`bool`): whether single_step forward, if so, add timestep dim before forward and\
remove it after forward
Return:
- ret (:obj:`dict`): output data dict with keys ['total_q', 'logit', 'next_state']
- total_q (:obj:`torch.Tensor`): total q_value, which is the result of mixer network
- agent_q (:obj:`torch.Tensor`): each agent q_value
- next_state (:obj:`list`): next rnn state
Shapes:
- agent_state (:obj:`torch.Tensor`): :math:`(T, B, A, N)`, where T is timestep, B is batch_size\
A is agent_num, N is obs_shape
- global_state (:obj:`torch.Tensor`): :math:`(T, B, M)`, where M is global_obs_shape
- prev_state (:obj:`list`): math:`(B, A)`, a list of length B, and each element is a list of length A
- action (:obj:`torch.Tensor`): :math:`(T, B, A)`
- total_q (:obj:`torch.Tensor`): :math:`(T, B)`
- agent_q (:obj:`torch.Tensor`): :math:`(T, B, A, P)`, where P is action_shape
- next_state (:obj:`list`): math:`(B, A)`, a list of length B, and each element is a list of length A
Examples:
>>> collaQ_model = CollaQ(
>>> agent_num=4,
>>> obs_shape=32,
>>> alone_obs_shape=24,
>>> global_obs_shape=32 * 4,
>>> action_shape=9,
>>> hidden_size_list=[128, 64],
>>> self_feature_range=[8, 10],
>>> ally_feature_range=[10, 16],
>>> attention_size=64,
>>> mixer=True,
>>> activation=torch.nn.Tanh()
>>> )
>>> data={
>>> 'obs': {
>>> 'agent_state': torch.randn(8, 4, 4, 32),
>>> 'agent_alone_state': torch.randn(8, 4, 4, 24),
>>> 'agent_alone_padding_state': torch.randn(8, 4, 4, 32),
>>> 'global_state': torch.randn(8, 4, 32 * 4),
>>> 'action_mask': torch.randint(0, 2, size=(8, 4, 4, 9))
>>> },
>>> 'prev_state': [[[None for _ in range(4)] for _ in range(3)] for _ in range(4)],
>>> 'action': torch.randint(0, 9, size=(8, 4, 4))
>>> }
>>> output = collaQ_model(data, single_step=False)
"""
agent_state, agent_alone_state = data['obs']['agent_state'], data['obs']['agent_alone_state']
agent_alone_padding_state = data['obs']['agent_alone_padding_state']
global_state, prev_state = data['obs']['global_state'], data['prev_state']
# TODO find a better way to implement agent_along_padding_state
action = data.get('action', None)
if single_step:
agent_state, agent_alone_state, agent_alone_padding_state, global_state = agent_state.unsqueeze(
0
), agent_alone_state.unsqueeze(0), agent_alone_padding_state.unsqueeze(0), global_state.unsqueeze(0)
T, B, A = agent_state.shape[:3]
if self.attention:
agent_state = self._self_attention(agent_state)
agent_alone_padding_state = self._self_attention(agent_alone_padding_state)
# prev state should be of size (B, 3, A) hidden_size)
"""
Note: to achieve such work, we should change the init_fn of hidden_state plugin in collaQ policy
"""
assert len(prev_state) == B and all([len(p) == 3 for p in prev_state]) and all(
[len(q) == A] for p in prev_state for q in p
), '{}-{}-{}-{}'.format([type(p) for p in prev_state], B, A, len(prev_state[0]))
alone_prev_state = [[None for _ in range(A)] for _ in range(B)]
colla_prev_state = [[None for _ in range(A)] for _ in range(B)]
colla_alone_prev_state = [[None for _ in range(A)] for _ in range(B)]
for i in range(B):
for j in range(3):
for k in range(A):
if j == 0:
alone_prev_state[i][k] = prev_state[i][j][k]
elif j == 1:
colla_prev_state[i][k] = prev_state[i][j][k]
elif j == 2:
colla_alone_prev_state[i][k] = prev_state[i][j][k]
alone_prev_state = reduce(lambda x, y: x + y, alone_prev_state)
colla_prev_state = reduce(lambda x, y: x + y, colla_prev_state)
colla_alone_prev_state = reduce(lambda x, y: x + y, colla_alone_prev_state)
agent_state = agent_state.reshape(T, -1, *agent_state.shape[3:])
agent_alone_state = agent_alone_state.reshape(T, -1, *agent_alone_state.shape[3:])
agent_alone_padding_state = agent_alone_padding_state.reshape(T, -1, *agent_alone_padding_state.shape[3:])
colla_output = self._q_network(
{
'obs': agent_state,
'prev_state': colla_prev_state,
'enable_fast_timestep': True
}
)
colla_alone_output = self._q_network(
{
'obs': agent_alone_padding_state,
'prev_state': colla_alone_prev_state,
'enable_fast_timestep': True
}
)
alone_output = self._q_alone_network(
{
'obs': agent_alone_state,
'prev_state': alone_prev_state,
'enable_fast_timestep': True
}
)
agent_alone_q, alone_next_state = alone_output['logit'], alone_output['next_state']
agent_colla_alone_q, colla_alone_next_state = colla_alone_output['logit'], colla_alone_output['next_state']
agent_colla_q, colla_next_state = colla_output['logit'], colla_output['next_state']
colla_next_state, _ = list_split(colla_next_state, step=A)
alone_next_state, _ = list_split(alone_next_state, step=A)
colla_alone_next_state, _ = list_split(colla_alone_next_state, step=A)
next_state = list(
map(lambda x: [x[0], x[1], x[2]], zip(alone_next_state, colla_next_state, colla_alone_next_state))
)
agent_alone_q = agent_alone_q.reshape(T, B, A, -1)
agent_colla_alone_q = agent_colla_alone_q.reshape(T, B, A, -1)
agent_colla_q = agent_colla_q.reshape(T, B, A, -1)
total_q_before_mix = agent_alone_q + agent_colla_q - agent_colla_alone_q
# total_q_before_mix = agent_colla_q
# total_q_before_mix = agent_alone_q
agent_q = total_q_before_mix
if action is None:
# For target forward process
if len(data['obs']['action_mask'].shape) == 3:
action_mask = data['obs']['action_mask'].unsqueeze(0)
else:
action_mask = data['obs']['action_mask']
agent_q[action_mask == 0.0] = -9999999
action = agent_q.argmax(dim=-1)
agent_q_act = torch.gather(agent_q, dim=-1, index=action.unsqueeze(-1))
agent_q_act = agent_q_act.squeeze(-1) # T, B, A
if self.mixer:
global_state_embedding = self._global_state_encoder(global_state)
total_q = self._mixer(agent_q_act, global_state_embedding)
else:
total_q = agent_q_act.sum(-1)
if single_step:
total_q, agent_q, agent_colla_alone_q = total_q.squeeze(0), agent_q.squeeze(0), agent_colla_alone_q.squeeze(
0
)
return {
'total_q': total_q,
'logit': agent_q,
'agent_colla_alone_q': agent_colla_alone_q * data['obs']['action_mask'],
'next_state': next_state,
'action_mask': data['obs']['action_mask']
}
def _setup_global_encoder(self, global_obs_shape: int, embedding_size: int) -> torch.nn.Module:
"""
Overview:
Used to encoder global observation.
Arguments:
- global_obs_shape (:obj:`int`): the dimension of global observation state
- embedding_size (:obj:`int`): the dimension of state emdedding
Returns:
- outputs (:obj:`torch.nn.Module`): Global observation encoding network
"""
return MLP(global_obs_shape, embedding_size, embedding_size, 2, activation=self._act)