Skip to content

ding.model.template.pdqn

ding.model.template.pdqn

PDQN

Bases: 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.

__init__(obs_shape, action_shape, encoder_hidden_size_list=[128, 128, 64], dueling=True, head_hidden_size=None, head_layer_num=1, activation=nn.ReLU(), norm_type=None, multi_pass=False, action_mask=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.

forward(inputs, mode)

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.

compute_continuous(inputs)

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])

compute_discrete(inputs)

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])

Full Source Code

../ding/model/template/pdqn.py

1from typing import Union, Optional, Dict 2from easydict import EasyDict 3 4import torch 5import torch.nn as nn 6 7from ding.torch_utils import get_lstm 8from ding.utils import MODEL_REGISTRY, SequenceType, squeeze 9from ..common import FCEncoder, ConvEncoder, DiscreteHead, DuelingHead, RegressionHead 10 11 12@MODEL_REGISTRY.register('pdqn') 13class PDQN(nn.Module): 14 """ 15 Overview: 16 The neural network and computation graph of PDQN(https://arxiv.org/abs/1810.06394v1) and \ 17 MPDQN(https://arxiv.org/abs/1905.04388) algorithms for parameterized action space. \ 18 This model supports parameterized action space with discrete ``action_type`` and continuous ``action_arg``. \ 19 In principle, PDQN consists of x network (continuous action parameter network) and Q network (discrete \ 20 action type network). But for simplicity, the code is split into ``encoder`` and ``actor_head``, which \ 21 contain the encoder and head of the above two networks respectively. 22 Interface: 23 ``__init__``, ``forward``, ``compute_discrete``, ``compute_continuous``. 24 """ 25 mode = ['compute_discrete', 'compute_continuous'] 26 27 def __init__( 28 self, 29 obs_shape: Union[int, SequenceType], 30 action_shape: EasyDict, 31 encoder_hidden_size_list: SequenceType = [128, 128, 64], 32 dueling: bool = True, 33 head_hidden_size: Optional[int] = None, 34 head_layer_num: int = 1, 35 activation: Optional[nn.Module] = nn.ReLU(), 36 norm_type: Optional[str] = None, 37 multi_pass: Optional[bool] = False, 38 action_mask: Optional[list] = None 39 ) -> None: 40 """ 41 Overview: 42 Init the PDQN (encoder + head) Model according to input arguments. 43 Arguments: 44 - obs_shape (:obj:`Union[int, SequenceType]`): Observation space shape, such as 8 or [4, 84, 84]. 45 - action_shape (:obj:`EasyDict`): Action space shape in dict type, such as \ 46 EasyDict({'action_type_shape': 3, 'action_args_shape': 5}). 47 - encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder``, \ 48 the last element must match ``head_hidden_size``. 49 - dueling (:obj:`dueling`): Whether choose ``DuelingHead`` or ``DiscreteHead(default)``. 50 - head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` of head network. 51 - head_layer_num (:obj:`int`): The number of layers used in the head network to compute Q value output. 52 - activation (:obj:`Optional[nn.Module]`): The type of activation function in networks \ 53 if ``None`` then default set it to ``nn.ReLU()``. 54 - norm_type (:obj:`Optional[str]`): The type of normalization in networks, see \ 55 ``ding.torch_utils.fc_block`` for more details. 56 - multi_pass (:obj:`Optional[bool]`): Whether to use multi pass version. 57 - action_mask: (:obj:`Optional[list]`): An action mask indicating how action args are \ 58 associated to each discrete action. For example, if there are 3 discrete action, \ 59 4 continous action args, and the first discrete action associates with the first \ 60 continuous action args, the second discrete action associates with the second continuous \ 61 action args, and the third discrete action associates with the remaining 2 action args, \ 62 the action mask will be like: [[1,0,0,0],[0,1,0,0],[0,0,1,1]] with shape 3*4. 63 """ 64 super(PDQN, self).__init__() 65 self.multi_pass = multi_pass 66 if self.multi_pass: 67 assert isinstance( 68 action_mask, list 69 ), 'Please indicate action mask in list form if you set multi_pass to True' 70 self.action_mask = torch.LongTensor(action_mask) 71 nonzero = torch.nonzero(self.action_mask) 72 index = torch.zeros(action_shape.action_args_shape).long() 73 index.scatter_(dim=0, index=nonzero[:, 1], src=nonzero[:, 0]) 74 self.action_scatter_index = index # (self.action_args_shape, ) 75 76 # squeeze action shape input like (3,) to 3 77 action_shape.action_args_shape = squeeze(action_shape.action_args_shape) 78 action_shape.action_type_shape = squeeze(action_shape.action_type_shape) 79 self.action_args_shape = action_shape.action_args_shape 80 self.action_type_shape = action_shape.action_type_shape 81 82 # init head hidden size 83 if head_hidden_size is None: 84 head_hidden_size = encoder_hidden_size_list[-1] 85 86 # squeeze obs input for compatibility: 1, (1, ), [4, 32, 32] 87 obs_shape = squeeze(obs_shape) 88 89 # Obs Encoder Type 90 if isinstance(obs_shape, int) or len(obs_shape) == 1: # FC Encoder 91 self.dis_encoder = FCEncoder( 92 obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type 93 ) 94 self.cont_encoder = FCEncoder( 95 obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type 96 ) 97 elif len(obs_shape) == 3: # Conv Encoder 98 self.dis_encoder = ConvEncoder( 99 obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type 100 ) 101 self.cont_encoder = ConvEncoder( 102 obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type 103 ) 104 else: 105 raise RuntimeError( 106 "Pre-defined encoder not support obs_shape {}, please customize your own PDQN.".format(obs_shape) 107 ) 108 109 # Continuous Action Head Type 110 self.cont_head = RegressionHead( 111 head_hidden_size, 112 action_shape.action_args_shape, 113 head_layer_num, 114 final_tanh=True, 115 activation=activation, 116 norm_type=norm_type 117 ) 118 119 # Discrete Action Head Type 120 if dueling: 121 dis_head_cls = DuelingHead 122 else: 123 dis_head_cls = DiscreteHead 124 self.dis_head = dis_head_cls( 125 head_hidden_size + action_shape.action_args_shape, 126 action_shape.action_type_shape, 127 head_layer_num, 128 activation=activation, 129 norm_type=norm_type 130 ) 131 132 self.actor_head = nn.ModuleList([self.dis_head, self.cont_head]) 133 # self.encoder = nn.ModuleList([self.dis_encoder, self.cont_encoder]) 134 # To speed up the training process, the X network and the Q network share the encoder for the state 135 self.encoder = nn.ModuleList([self.cont_encoder, self.cont_encoder]) 136 137 def forward(self, inputs: Union[torch.Tensor, Dict, EasyDict], mode: str) -> Dict: 138 """ 139 Overview: 140 PDQN forward computation graph, input observation tensor to predict q_value for \ 141 discrete actions and values for continuous action_args. 142 Arguments: 143 - inputs (:obj:`Union[torch.Tensor, Dict, EasyDict]`): Inputs including observation and \ 144 other info according to `mode`. 145 - mode (:obj:`str`): Name of the forward mode. 146 Shapes: 147 - inputs (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is ``obs_shape``. 148 """ 149 assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode) 150 return getattr(self, mode)(inputs) 151 152 def compute_continuous(self, inputs: torch.Tensor) -> Dict: 153 """ 154 Overview: 155 Use observation tensor to predict continuous action args. 156 Arguments: 157 - inputs (:obj:`torch.Tensor`): Observation inputs. 158 Returns: 159 - outputs (:obj:`Dict`): A dict with key 'action_args'. 160 - 'action_args' (:obj:`torch.Tensor`): The continuous action args. 161 Shapes: 162 - inputs (:obj:`torch.Tensor`): :math:`(B, N)`, where B is batch size and N is ``obs_shape``. 163 - action_args (:obj:`torch.Tensor`): :math:`(B, M)`, where M is ``action_args_shape``. 164 Examples: 165 >>> act_shape = EasyDict({'action_type_shape': (3, ), 'action_args_shape': (5, )}) 166 >>> model = PDQN(4, act_shape) 167 >>> inputs = torch.randn(64, 4) 168 >>> outputs = model.forward(inputs, mode='compute_continuous') 169 >>> assert outputs['action_args'].shape == torch.Size([64, 5]) 170 """ 171 cont_x = self.encoder[1](inputs) # size (B, encoded_state_shape) 172 action_args = self.actor_head[1](cont_x)['pred'] # size (B, action_args_shape) 173 outputs = {'action_args': action_args} 174 return outputs 175 176 def compute_discrete(self, inputs: Union[Dict, EasyDict]) -> Dict: 177 """ 178 Overview: 179 Use observation tensor and continuous action args to predict discrete action types. 180 Arguments: 181 - inputs (:obj:`Union[Dict, EasyDict]`): A dict with keys 'state', 'action_args'. 182 - state (:obj:`torch.Tensor`): Observation inputs. 183 - action_args (:obj:`torch.Tensor`): Action parameters are used to concatenate with the observation \ 184 and serve as input to the discrete action type network. 185 Returns: 186 - outputs (:obj:`Dict`): A dict with keys 'logit', 'action_args'. 187 - 'logit': The logit value for each discrete action. 188 - 'action_args': The continuous action args(same as the inputs['action_args']) for later usage. 189 Examples: 190 >>> act_shape = EasyDict({'action_type_shape': (3, ), 'action_args_shape': (5, )}) 191 >>> model = PDQN(4, act_shape) 192 >>> inputs = {'state': torch.randn(64, 4), 'action_args': torch.randn(64, 5)} 193 >>> outputs = model.forward(inputs, mode='compute_discrete') 194 >>> assert outputs['logit'].shape == torch.Size([64, 3]) 195 >>> assert outputs['action_args'].shape == torch.Size([64, 5]) 196 """ 197 dis_x = self.encoder[0](inputs['state']) # size (B, encoded_state_shape) 198 action_args = inputs['action_args'] # size (B, action_args_shape) 199 200 if self.multi_pass: # mpdqn 201 # fill_value=-2 is a mask value, which is not in normal acton range 202 # (B, action_args_shape, K) where K is the action_type_shape 203 mp_action = torch.full( 204 (dis_x.shape[0], self.action_args_shape, self.action_type_shape), 205 fill_value=-2, 206 device=dis_x.device, 207 dtype=dis_x.dtype 208 ) 209 index = self.action_scatter_index.view(1, -1, 1).repeat(dis_x.shape[0], 1, 1).to(dis_x.device) 210 211 # index: (B, action_args_shape, 1) src: (B, action_args_shape, 1) 212 mp_action.scatter_(dim=-1, index=index, src=action_args.unsqueeze(-1)) 213 mp_action = mp_action.permute(0, 2, 1) # (B, K, action_args_shape) 214 215 mp_state = dis_x.unsqueeze(1).repeat(1, self.action_type_shape, 1) # (B, K, obs_shape) 216 mp_state_action_cat = torch.cat([mp_state, mp_action], dim=-1) 217 218 logit = self.actor_head[0](mp_state_action_cat)['logit'] # (B, K, K) 219 220 logit = torch.diagonal(logit, dim1=-2, dim2=-1) # (B, K) 221 else: # pdqn 222 # size (B, encoded_state_shape + action_args_shape) 223 if len(action_args.shape) == 1: # (B, ) -> (B, 1) 224 action_args = action_args.unsqueeze(1) 225 state_action_cat = torch.cat((dis_x, action_args), dim=-1) 226 logit = self.actor_head[0](state_action_cat)['logit'] # size (B, K) where K is action_type_shape 227 228 outputs = {'logit': logit, 'action_args': action_args} 229 return outputs