Source code for ding.model.template.maqac
from typing import Union, Dict, Optional
from easydict import EasyDict
import torch
import torch.nn as nn
from ding.utils import SequenceType, squeeze, MODEL_REGISTRY
from ..common import RegressionHead, ReparameterizationHead, DiscreteHead, MultiHead, \
FCEncoder, ConvEncoder
[docs]@MODEL_REGISTRY.register('discrete_maqac')
class DiscreteMAQAC(nn.Module):
"""
Overview:
The neural network and computation graph of algorithms related to discrete action Multi-Agent Q-value \
Actor-CritiC (MAQAC) model. The model is composed of actor and critic, where actor is a MLP network and \
critic is a MLP network. The actor network is used to predict the action probability distribution, and the \
critic network is used to predict the Q value of the state-action pair.
Interfaces:
``__init__``, ``forward``, ``compute_actor``, ``compute_critic``
"""
mode = ['compute_actor', 'compute_critic']
[docs] def __init__(
self,
agent_obs_shape: Union[int, SequenceType],
global_obs_shape: Union[int, SequenceType],
action_shape: Union[int, SequenceType],
twin_critic: bool = False,
actor_head_hidden_size: int = 64,
actor_head_layer_num: int = 1,
critic_head_hidden_size: int = 64,
critic_head_layer_num: int = 1,
activation: Optional[nn.Module] = nn.ReLU(),
norm_type: Optional[str] = None,
) -> None:
"""
Overview:
Initialize the DiscreteMAQAC Model according to arguments.
Arguments:
- agent_obs_shape (:obj:`Union[int, SequenceType]`): Agent's observation's space.
- global_obs_shape (:obj:`Union[int, SequenceType]`): Global observation's space.
- obs_shape (:obj:`Union[int, SequenceType]`): Observation's space.
- action_shape (:obj:`Union[int, SequenceType]`): Action's space.
- twin_critic (:obj:`bool`): Whether include twin critic.
- actor_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to actor-nn's ``Head``.
- actor_head_layer_num (:obj:`int`): The num of layers used in the network to compute Q value output \
for actor's nn.
- critic_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to critic-nn's ``Head``.
- critic_head_layer_num (:obj:`int`): The num of layers used in the network to compute Q value output \
for critic's nn.
- activation (:obj:`Optional[nn.Module]`): The type of activation function to use in ``MLP`` the after \
``layer_fn``, if ``None`` then default set to ``nn.ReLU()``
- norm_type (:obj:`Optional[str]`): The type of normalization to use, see ``ding.torch_utils.fc_block`` \
for more details.
"""
super(DiscreteMAQAC, self).__init__()
agent_obs_shape: int = squeeze(agent_obs_shape)
action_shape: int = squeeze(action_shape)
self.actor = nn.Sequential(
nn.Linear(agent_obs_shape, actor_head_hidden_size), activation,
DiscreteHead(
actor_head_hidden_size, action_shape, actor_head_layer_num, activation=activation, norm_type=norm_type
)
)
self.twin_critic = twin_critic
if self.twin_critic:
self.critic = nn.ModuleList()
for _ in range(2):
self.critic.append(
nn.Sequential(
nn.Linear(global_obs_shape, critic_head_hidden_size), activation,
DiscreteHead(
critic_head_hidden_size,
action_shape,
critic_head_layer_num,
activation=activation,
norm_type=norm_type
)
)
)
else:
self.critic = nn.Sequential(
nn.Linear(global_obs_shape, critic_head_hidden_size), activation,
DiscreteHead(
critic_head_hidden_size,
action_shape,
critic_head_layer_num,
activation=activation,
norm_type=norm_type
)
)
[docs] def forward(self, inputs: Union[torch.Tensor, Dict], mode: str) -> Dict:
"""
Overview:
Use observation tensor to predict output, with ``compute_actor`` or ``compute_critic`` mode.
Arguments:
- inputs (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys:
- ``obs`` (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys:
- ``agent_state`` (:obj:`torch.Tensor`): The agent's observation tensor data, \
with shape :math:`(B, A, N0)`, where B is batch size and A is agent num. \
N0 corresponds to ``agent_obs_shape``.
- ``global_state`` (:obj:`torch.Tensor`): The global observation tensor data, \
with shape :math:`(B, A, N1)`, where B is batch size and A is agent num. \
N1 corresponds to ``global_obs_shape``.
- ``action_mask`` (:obj:`torch.Tensor`): The action mask tensor data, \
with shape :math:`(B, A, N2)`, where B is batch size and A is agent num. \
N2 corresponds to ``action_shape``.
- mode (:obj:`str`): The forward mode, all the modes are defined in the beginning of this class.
Returns:
- output (:obj:`Dict[str, torch.Tensor]`): The output dict of DiscreteMAQAC forward computation graph, \
whose key-values vary in different forward modes.
Examples:
>>> B = 32
>>> agent_obs_shape = 216
>>> global_obs_shape = 264
>>> agent_num = 8
>>> action_shape = 14
>>> data = {
>>> 'obs': {
>>> 'agent_state': torch.randn(B, agent_num, agent_obs_shape),
>>> 'global_state': torch.randn(B, agent_num, global_obs_shape),
>>> 'action_mask': torch.randint(0, 2, size=(B, agent_num, action_shape))
>>> }
>>> }
>>> model = DiscreteMAQAC(agent_obs_shape, global_obs_shape, action_shape, twin_critic=True)
>>> logit = model(data, mode='compute_actor')['logit']
>>> value = model(data, mode='compute_critic')['q_value']
"""
assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode)
return getattr(self, mode)(inputs)
[docs] def compute_actor(self, inputs: Dict) -> Dict:
"""
Overview:
Use observation tensor to predict action logits.
Arguments:
- inputs (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys:
- ``obs`` (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys:
- ``agent_state`` (:obj:`torch.Tensor`): The agent's observation tensor data, \
with shape :math:`(B, A, N0)`, where B is batch size and A is agent num. \
N0 corresponds to ``agent_obs_shape``.
- ``global_state`` (:obj:`torch.Tensor`): The global observation tensor data, \
with shape :math:`(B, A, N1)`, where B is batch size and A is agent num. \
N1 corresponds to ``global_obs_shape``.
- ``action_mask`` (:obj:`torch.Tensor`): The action mask tensor data, \
with shape :math:`(B, A, N2)`, where B is batch size and A is agent num. \
N2 corresponds to ``action_shape``.
Returns:
- output (:obj:`Dict[str, torch.Tensor]`): The output dict of DiscreteMAQAC forward computation graph, \
whose key-values vary in different forward modes.
- logit (:obj:`torch.Tensor`): Action's output logit (real value range), whose shape is \
:math:`(B, A, N2)`, where N2 corresponds to ``action_shape``.
- action_mask (:obj:`torch.Tensor`): Action mask tensor with same size as ``action_shape``.
Examples:
>>> B = 32
>>> agent_obs_shape = 216
>>> global_obs_shape = 264
>>> agent_num = 8
>>> action_shape = 14
>>> data = {
>>> 'obs': {
>>> 'agent_state': torch.randn(B, agent_num, agent_obs_shape),
>>> 'global_state': torch.randn(B, agent_num, global_obs_shape),
>>> 'action_mask': torch.randint(0, 2, size=(B, agent_num, action_shape))
>>> }
>>> }
>>> model = DiscreteMAQAC(agent_obs_shape, global_obs_shape, action_shape, twin_critic=True)
>>> logit = model.compute_actor(data)['logit']
"""
action_mask = inputs['obs']['action_mask']
x = self.actor(inputs['obs']['agent_state'])
return {'logit': x['logit'], 'action_mask': action_mask}
[docs] def compute_critic(self, inputs: Dict) -> Dict:
"""
Overview:
use observation tensor to predict Q value.
Arguments:
- inputs (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys:
- ``obs`` (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys:
- ``agent_state`` (:obj:`torch.Tensor`): The agent's observation tensor data, \
with shape :math:`(B, A, N0)`, where B is batch size and A is agent num. \
N0 corresponds to ``agent_obs_shape``.
- ``global_state`` (:obj:`torch.Tensor`): The global observation tensor data, \
with shape :math:`(B, A, N1)`, where B is batch size and A is agent num. \
N1 corresponds to ``global_obs_shape``.
- ``action_mask`` (:obj:`torch.Tensor`): The action mask tensor data, \
with shape :math:`(B, A, N2)`, where B is batch size and A is agent num. \
N2 corresponds to ``action_shape``.
Returns:
- output (:obj:`Dict[str, torch.Tensor]`): The output dict of DiscreteMAQAC forward computation graph, \
whose key-values vary in different values of ``twin_critic``.
- q_value (:obj:`list`): If ``twin_critic=True``, q_value should be 2 elements, each is the shape of \
:math:`(B, A, N2)`, where B is batch size and A is agent num. N2 corresponds to ``action_shape``. \
Otherwise, q_value should be ``torch.Tensor``.
Examples:
>>> B = 32
>>> agent_obs_shape = 216
>>> global_obs_shape = 264
>>> agent_num = 8
>>> action_shape = 14
>>> data = {
>>> 'obs': {
>>> 'agent_state': torch.randn(B, agent_num, agent_obs_shape),
>>> 'global_state': torch.randn(B, agent_num, global_obs_shape),
>>> 'action_mask': torch.randint(0, 2, size=(B, agent_num, action_shape))
>>> }
>>> }
>>> model = DiscreteMAQAC(agent_obs_shape, global_obs_shape, action_shape, twin_critic=True)
>>> value = model.compute_critic(data)['q_value']
"""
if self.twin_critic:
x = [m(inputs['obs']['global_state'])['logit'] for m in self.critic]
else:
x = self.critic(inputs['obs']['global_state'])['logit']
return {'q_value': x}
[docs]@MODEL_REGISTRY.register('continuous_maqac')
class ContinuousMAQAC(nn.Module):
"""
Overview:
The neural network and computation graph of algorithms related to continuous action Multi-Agent Q-value \
Actor-CritiC (MAQAC) model. The model is composed of actor and critic, where actor is a MLP network and \
critic is a MLP network. The actor network is used to predict the action probability distribution, and the \
critic network is used to predict the Q value of the state-action pair.
Interfaces:
``__init__``, ``forward``, ``compute_actor``, ``compute_critic``
"""
mode = ['compute_actor', 'compute_critic']
[docs] def __init__(
self,
agent_obs_shape: Union[int, SequenceType],
global_obs_shape: Union[int, SequenceType],
action_shape: Union[int, SequenceType, EasyDict],
action_space: str,
twin_critic: bool = False,
actor_head_hidden_size: int = 64,
actor_head_layer_num: int = 1,
critic_head_hidden_size: int = 64,
critic_head_layer_num: int = 1,
activation: Optional[nn.Module] = nn.ReLU(),
norm_type: Optional[str] = None,
) -> None:
"""
Overview:
Initialize the QAC Model according to arguments.
Arguments:
- obs_shape (:obj:`Union[int, SequenceType]`): Observation's space.
- action_shape (:obj:`Union[int, SequenceType, EasyDict]`): Action's space, such as 4, (3, )
- action_space (:obj:`str`): Whether choose ``regression`` or ``reparameterization``.
- twin_critic (:obj:`bool`): Whether include twin critic.
- actor_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to actor-nn's ``Head``.
- actor_head_layer_num (:obj:`int`): The num of layers used in the network to compute Q value output \
for actor's nn.
- critic_head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to critic-nn's ``Head``.
- critic_head_layer_num (:obj:`int`): The num of layers used in the network to compute Q value output \
for critic's nn.
- activation (:obj:`Optional[nn.Module]`): The type of activation function to use in ``MLP`` the after \
``layer_fn``, if ``None`` then default set to ``nn.ReLU()``
- norm_type (:obj:`Optional[str]`): The type of normalization to use, see ``ding.torch_utils.fc_block`` \
for more details.
"""
super(ContinuousMAQAC, self).__init__()
obs_shape: int = squeeze(agent_obs_shape)
global_obs_shape: int = squeeze(global_obs_shape)
action_shape = squeeze(action_shape)
self.action_shape = action_shape
self.action_space = action_space
assert self.action_space in ['regression', 'reparameterization'], self.action_space
if self.action_space == 'regression': # DDPG, TD3
self.actor = nn.Sequential(
nn.Linear(obs_shape, actor_head_hidden_size), activation,
RegressionHead(
actor_head_hidden_size,
action_shape,
actor_head_layer_num,
final_tanh=True,
activation=activation,
norm_type=norm_type
)
)
else: # SAC
self.actor = nn.Sequential(
nn.Linear(obs_shape, actor_head_hidden_size), activation,
ReparameterizationHead(
actor_head_hidden_size,
action_shape,
actor_head_layer_num,
sigma_type='conditioned',
activation=activation,
norm_type=norm_type
)
)
self.twin_critic = twin_critic
critic_input_size = global_obs_shape + action_shape
if self.twin_critic:
self.critic = nn.ModuleList()
for _ in range(2):
self.critic.append(
nn.Sequential(
nn.Linear(critic_input_size, critic_head_hidden_size), activation,
RegressionHead(
critic_head_hidden_size,
1,
critic_head_layer_num,
final_tanh=False,
activation=activation,
norm_type=norm_type
)
)
)
else:
self.critic = nn.Sequential(
nn.Linear(critic_input_size, critic_head_hidden_size), activation,
RegressionHead(
critic_head_hidden_size,
1,
critic_head_layer_num,
final_tanh=False,
activation=activation,
norm_type=norm_type
)
)
[docs] def forward(self, inputs: Union[torch.Tensor, Dict], mode: str) -> Dict:
"""
Overview:
Use observation and action tensor to predict output in ``compute_actor`` or ``compute_critic`` mode.
Arguments:
- inputs (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys:
- ``obs`` (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys:
- ``agent_state`` (:obj:`torch.Tensor`): The agent's observation tensor data, \
with shape :math:`(B, A, N0)`, where B is batch size and A is agent num. \
N0 corresponds to ``agent_obs_shape``.
- ``global_state`` (:obj:`torch.Tensor`): The global observation tensor data, \
with shape :math:`(B, A, N1)`, where B is batch size and A is agent num. \
N1 corresponds to ``global_obs_shape``.
- ``action_mask`` (:obj:`torch.Tensor`): The action mask tensor data, \
with shape :math:`(B, A, N2)`, where B is batch size and A is agent num. \
N2 corresponds to ``action_shape``.
- ``action`` (:obj:`torch.Tensor`): The action tensor data, \
with shape :math:`(B, A, N3)`, where B is batch size and A is agent num. \
N3 corresponds to ``action_shape``.
- mode (:obj:`str`): Name of the forward mode.
Returns:
- outputs (:obj:`Dict`): Outputs of network forward, whose key-values will be different for different \
``mode``, ``twin_critic``, ``action_space``.
Examples:
>>> B = 32
>>> agent_obs_shape = 216
>>> global_obs_shape = 264
>>> agent_num = 8
>>> action_shape = 14
>>> act_space = 'reparameterization' # regression
>>> data = {
>>> 'obs': {
>>> 'agent_state': torch.randn(B, agent_num, agent_obs_shape),
>>> 'global_state': torch.randn(B, agent_num, global_obs_shape),
>>> 'action_mask': torch.randint(0, 2, size=(B, agent_num, action_shape))
>>> },
>>> 'action': torch.randn(B, agent_num, squeeze(action_shape))
>>> }
>>> model = ContinuousMAQAC(agent_obs_shape, global_obs_shape, action_shape, act_space, twin_critic=False)
>>> if action_space == 'regression':
>>> action = model(data['obs'], mode='compute_actor')['action']
>>> elif action_space == 'reparameterization':
>>> (mu, sigma) = model(data['obs'], mode='compute_actor')['logit']
>>> value = model(data, mode='compute_critic')['q_value']
"""
assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode)
return getattr(self, mode)(inputs)
[docs] def compute_actor(self, inputs: Dict) -> Dict:
"""
Overview:
Use observation tensor to predict action logits.
Arguments:
- inputs (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys:
- ``agent_state`` (:obj:`torch.Tensor`): The agent's observation tensor data, \
with shape :math:`(B, A, N0)`, where B is batch size and A is agent num. \
N0 corresponds to ``agent_obs_shape``.
Returns:
- outputs (:obj:`Dict`): Outputs of network forward.
ReturnKeys (``action_space == 'regression'``):
- action (:obj:`torch.Tensor`): Action tensor with same size as ``action_shape``.
ReturnKeys (``action_space == 'reparameterization'``):
- logit (:obj:`list`): 2 elements, each is the shape of :math:`(B, A, N3)`, where B is batch size and \
A is agent num. N3 corresponds to ``action_shape``.
Examples:
>>> B = 32
>>> agent_obs_shape = 216
>>> global_obs_shape = 264
>>> agent_num = 8
>>> action_shape = 14
>>> act_space = 'reparameterization' # 'regression'
>>> data = {
>>> 'agent_state': torch.randn(B, agent_num, agent_obs_shape),
>>> }
>>> model = ContinuousMAQAC(agent_obs_shape, global_obs_shape, action_shape, act_space, twin_critic=False)
>>> if action_space == 'regression':
>>> action = model.compute_actor(data)['action']
>>> elif action_space == 'reparameterization':
>>> (mu, sigma) = model.compute_actor(data)['logit']
"""
inputs = inputs['agent_state']
if self.action_space == 'regression':
x = self.actor(inputs)
return {'action': x['pred']}
else:
x = self.actor(inputs)
return {'logit': [x['mu'], x['sigma']]}
[docs] def compute_critic(self, inputs: Dict) -> Dict:
"""
Overview:
Use observation tensor and action tensor to predict Q value.
Arguments:
- inputs (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys:
- ``obs`` (:obj:`Dict[str, torch.Tensor]`): The input dict tensor data, has keys:
- ``agent_state`` (:obj:`torch.Tensor`): The agent's observation tensor data, \
with shape :math:`(B, A, N0)`, where B is batch size and A is agent num. \
N0 corresponds to ``agent_obs_shape``.
- ``global_state`` (:obj:`torch.Tensor`): The global observation tensor data, \
with shape :math:`(B, A, N1)`, where B is batch size and A is agent num. \
N1 corresponds to ``global_obs_shape``.
- ``action_mask`` (:obj:`torch.Tensor`): The action mask tensor data, \
with shape :math:`(B, A, N2)`, where B is batch size and A is agent num. \
N2 corresponds to ``action_shape``.
- ``action`` (:obj:`torch.Tensor`): The action tensor data, \
with shape :math:`(B, A, N3)`, where B is batch size and A is agent num. \
N3 corresponds to ``action_shape``.
Returns:
- outputs (:obj:`Dict`): Outputs of network forward.
ReturnKeys (``twin_critic=True``):
- q_value (:obj:`list`): 2 elements, each is the shape of :math:`(B, A)`, where B is batch size and \
A is agent num.
ReturnKeys (``twin_critic=False``):
- q_value (:obj:`torch.Tensor`): :math:`(B, A)`, where B is batch size and A is agent num.
Examples:
>>> B = 32
>>> agent_obs_shape = 216
>>> global_obs_shape = 264
>>> agent_num = 8
>>> action_shape = 14
>>> act_space = 'reparameterization' # 'regression'
>>> data = {
>>> 'obs': {
>>> 'agent_state': torch.randn(B, agent_num, agent_obs_shape),
>>> 'global_state': torch.randn(B, agent_num, global_obs_shape),
>>> 'action_mask': torch.randint(0, 2, size=(B, agent_num, action_shape))
>>> },
>>> 'action': torch.randn(B, agent_num, squeeze(action_shape))
>>> }
>>> model = ContinuousMAQAC(agent_obs_shape, global_obs_shape, action_shape, act_space, twin_critic=False)
>>> value = model.compute_critic(data)['q_value']
"""
obs, action = inputs['obs']['global_state'], inputs['action']
if len(action.shape) == 1: # (B, ) -> (B, 1)
action = action.unsqueeze(1)
x = torch.cat([obs, action], dim=-1)
if self.twin_critic:
x = [m(x)['pred'] for m in self.critic]
else:
x = self.critic(x)['pred']
return {'q_value': x}