File size: 11,800 Bytes
079c32c |
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 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 |
from typing import Union, Optional, Dict
from easydict import EasyDict
import torch
import torch.nn as nn
from ding.torch_utils import get_lstm
from ding.utils import MODEL_REGISTRY, SequenceType, squeeze
from ..common import FCEncoder, ConvEncoder, DiscreteHead, DuelingHead, RegressionHead
@MODEL_REGISTRY.register('pdqn')
class PDQN(nn.Module):
"""
Overview:
The neural network and computation graph of PDQN(https://arxiv.org/abs/1810.06394v1) and \
MPDQN(https://arxiv.org/abs/1905.04388) algorithms for parameterized action space. \
This model supports parameterized action space with discrete ``action_type`` and continuous ``action_arg``. \
In principle, PDQN consists of x network (continuous action parameter network) and Q network (discrete \
action type network). But for simplicity, the code is split into ``encoder`` and ``actor_head``, which \
contain the encoder and head of the above two networks respectively.
Interface:
``__init__``, ``forward``, ``compute_discrete``, ``compute_continuous``.
"""
mode = ['compute_discrete', 'compute_continuous']
def __init__(
self,
obs_shape: Union[int, SequenceType],
action_shape: EasyDict,
encoder_hidden_size_list: SequenceType = [128, 128, 64],
dueling: bool = True,
head_hidden_size: Optional[int] = None,
head_layer_num: int = 1,
activation: Optional[nn.Module] = nn.ReLU(),
norm_type: Optional[str] = None,
multi_pass: Optional[bool] = False,
action_mask: Optional[list] = None
) -> None:
"""
Overview:
Init the PDQN (encoder + head) Model according to input arguments.
Arguments:
- obs_shape (:obj:`Union[int, SequenceType]`): Observation space shape, such as 8 or [4, 84, 84].
- action_shape (:obj:`EasyDict`): Action space shape in dict type, such as \
EasyDict({'action_type_shape': 3, 'action_args_shape': 5}).
- encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder``, \
the last element must match ``head_hidden_size``.
- dueling (:obj:`dueling`): Whether choose ``DuelingHead`` or ``DiscreteHead(default)``.
- head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` of head network.
- head_layer_num (:obj:`int`): The number of layers used in the head network to compute Q value output.
- activation (:obj:`Optional[nn.Module]`): The type of activation function in networks \
if ``None`` then default set it to ``nn.ReLU()``.
- norm_type (:obj:`Optional[str]`): The type of normalization in networks, see \
``ding.torch_utils.fc_block`` for more details.
- multi_pass (:obj:`Optional[bool]`): Whether to use multi pass version.
- action_mask: (:obj:`Optional[list]`): An action mask indicating how action args are \
associated to each discrete action. For example, if there are 3 discrete action, \
4 continous action args, and the first discrete action associates with the first \
continuous action args, the second discrete action associates with the second continuous \
action args, and the third discrete action associates with the remaining 2 action args, \
the action mask will be like: [[1,0,0,0],[0,1,0,0],[0,0,1,1]] with shape 3*4.
"""
super(PDQN, self).__init__()
self.multi_pass = multi_pass
if self.multi_pass:
assert isinstance(
action_mask, list
), 'Please indicate action mask in list form if you set multi_pass to True'
self.action_mask = torch.LongTensor(action_mask)
nonzero = torch.nonzero(self.action_mask)
index = torch.zeros(action_shape.action_args_shape).long()
index.scatter_(dim=0, index=nonzero[:, 1], src=nonzero[:, 0])
self.action_scatter_index = index # (self.action_args_shape, )
# squeeze action shape input like (3,) to 3
action_shape.action_args_shape = squeeze(action_shape.action_args_shape)
action_shape.action_type_shape = squeeze(action_shape.action_type_shape)
self.action_args_shape = action_shape.action_args_shape
self.action_type_shape = action_shape.action_type_shape
# init head hidden size
if head_hidden_size is None:
head_hidden_size = encoder_hidden_size_list[-1]
# squeeze obs input for compatibility: 1, (1, ), [4, 32, 32]
obs_shape = squeeze(obs_shape)
# Obs Encoder Type
if isinstance(obs_shape, int) or len(obs_shape) == 1: # FC Encoder
self.dis_encoder = FCEncoder(
obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type
)
self.cont_encoder = FCEncoder(
obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type
)
elif len(obs_shape) == 3: # Conv Encoder
self.dis_encoder = ConvEncoder(
obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type
)
self.cont_encoder = ConvEncoder(
obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type
)
else:
raise RuntimeError(
"Pre-defined encoder not support obs_shape {}, please customize your own PDQN.".format(obs_shape)
)
# Continuous Action Head Type
self.cont_head = RegressionHead(
head_hidden_size,
action_shape.action_args_shape,
head_layer_num,
final_tanh=True,
activation=activation,
norm_type=norm_type
)
# Discrete Action Head Type
if dueling:
dis_head_cls = DuelingHead
else:
dis_head_cls = DiscreteHead
self.dis_head = dis_head_cls(
head_hidden_size + action_shape.action_args_shape,
action_shape.action_type_shape,
head_layer_num,
activation=activation,
norm_type=norm_type
)
self.actor_head = nn.ModuleList([self.dis_head, self.cont_head])
# self.encoder = nn.ModuleList([self.dis_encoder, self.cont_encoder])
# To speed up the training process, the X network and the Q network share the encoder for the state
self.encoder = nn.ModuleList([self.cont_encoder, self.cont_encoder])
def forward(self, inputs: Union[torch.Tensor, Dict, EasyDict], mode: str) -> Dict:
"""
Overview:
PDQN forward computation graph, input observation tensor to predict q_value for \
discrete actions and values for continuous action_args.
Arguments:
- inputs (:obj:`Union[torch.Tensor, Dict, EasyDict]`): Inputs including observation and \
other info according to `mode`.
- mode (:obj:`str`): Name of the forward mode.
Shapes:
- inputs (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is ``obs_shape``.
"""
assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode)
return getattr(self, mode)(inputs)
def compute_continuous(self, inputs: torch.Tensor) -> Dict:
"""
Overview:
Use observation tensor to predict continuous action args.
Arguments:
- inputs (:obj:`torch.Tensor`): Observation inputs.
Returns:
- outputs (:obj:`Dict`): A dict with key 'action_args'.
- 'action_args' (:obj:`torch.Tensor`): The continuous action args.
Shapes:
- inputs (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is ``obs_shape``.
- action_args (:obj:`torch.Tensor`): :math:`(B, M)`, where M is ``action_args_shape``.
Examples:
>>> act_shape = EasyDict({'action_type_shape': (3, ), 'action_args_shape': (5, )})
>>> model = PDQN(4, act_shape)
>>> inputs = torch.randn(64, 4)
>>> outputs = model.forward(inputs, mode='compute_continuous')
>>> assert outputs['action_args'].shape == torch.Size([64, 5])
"""
cont_x = self.encoder[1](inputs) # size (B, encoded_state_shape)
action_args = self.actor_head[1](cont_x)['pred'] # size (B, action_args_shape)
outputs = {'action_args': action_args}
return outputs
def compute_discrete(self, inputs: Union[Dict, EasyDict]) -> Dict:
"""
Overview:
Use observation tensor and continuous action args to predict discrete action types.
Arguments:
- inputs (:obj:`Union[Dict, EasyDict]`): A dict with keys 'state', 'action_args'.
- state (:obj:`torch.Tensor`): Observation inputs.
- action_args (:obj:`torch.Tensor`): Action parameters are used to concatenate with the observation \
and serve as input to the discrete action type network.
Returns:
- outputs (:obj:`Dict`): A dict with keys 'logit', 'action_args'.
- 'logit': The logit value for each discrete action.
- 'action_args': The continuous action args(same as the inputs['action_args']) for later usage.
Examples:
>>> act_shape = EasyDict({'action_type_shape': (3, ), 'action_args_shape': (5, )})
>>> model = PDQN(4, act_shape)
>>> inputs = {'state': torch.randn(64, 4), 'action_args': torch.randn(64, 5)}
>>> outputs = model.forward(inputs, mode='compute_discrete')
>>> assert outputs['logit'].shape == torch.Size([64, 3])
>>> assert outputs['action_args'].shape == torch.Size([64, 5])
"""
dis_x = self.encoder[0](inputs['state']) # size (B, encoded_state_shape)
action_args = inputs['action_args'] # size (B, action_args_shape)
if self.multi_pass: # mpdqn
# fill_value=-2 is a mask value, which is not in normal acton range
# (B, action_args_shape, K) where K is the action_type_shape
mp_action = torch.full(
(dis_x.shape[0], self.action_args_shape, self.action_type_shape),
fill_value=-2,
device=dis_x.device,
dtype=dis_x.dtype
)
index = self.action_scatter_index.view(1, -1, 1).repeat(dis_x.shape[0], 1, 1).to(dis_x.device)
# index: (B, action_args_shape, 1) src: (B, action_args_shape, 1)
mp_action.scatter_(dim=-1, index=index, src=action_args.unsqueeze(-1))
mp_action = mp_action.permute(0, 2, 1) # (B, K, action_args_shape)
mp_state = dis_x.unsqueeze(1).repeat(1, self.action_type_shape, 1) # (B, K, obs_shape)
mp_state_action_cat = torch.cat([mp_state, mp_action], dim=-1)
logit = self.actor_head[0](mp_state_action_cat)['logit'] # (B, K, K)
logit = torch.diagonal(logit, dim1=-2, dim2=-1) # (B, K)
else: # pdqn
# size (B, encoded_state_shape + action_args_shape)
if len(action_args.shape) == 1: # (B, ) -> (B, 1)
action_args = action_args.unsqueeze(1)
state_action_cat = torch.cat((dis_x, action_args), dim=-1)
logit = self.actor_head[0](state_action_cat)['logit'] # size (B, K) where K is action_type_shape
outputs = {'logit': logit, 'action_args': action_args}
return outputs
|