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