Source code for ding.model.template.ngu
from typing import Union, Optional, Dict, Callable, List
import torch
import torch.nn as nn
from ding.torch_utils import get_lstm, one_hot, to_tensor, to_ndarray
from ding.utils import MODEL_REGISTRY, SequenceType, squeeze
# from ding.torch_utils.data_helper import one_hot_embedding, one_hot_embedding_none
from ..common import FCEncoder, ConvEncoder, DiscreteHead, DuelingHead, MultiHead, RainbowHead, \
QuantileHead, QRDQNHead, DistributionHead
def parallel_wrapper(forward_fn: Callable) -> Callable:
"""
Overview:
Process timestep T and batch_size B at the same time, in other words, treat different timestep data as \
different trajectories in a batch.
Arguments:
- forward_fn (:obj:`Callable`): Normal ``nn.Module`` 's forward function.
Returns:
- wrapper (:obj:`Callable`): Wrapped function.
"""
def wrapper(x: torch.Tensor) -> Union[torch.Tensor, List[torch.Tensor]]:
T, B = x.shape[:2]
def reshape(d):
if isinstance(d, list):
d = [reshape(t) for t in d]
elif isinstance(d, dict):
d = {k: reshape(v) for k, v in d.items()}
else:
d = d.reshape(T, B, *d.shape[1:])
return d
x = x.reshape(T * B, *x.shape[2:])
x = forward_fn(x)
x = reshape(x)
return x
return wrapper
[docs]@MODEL_REGISTRY.register('ngu')
class NGU(nn.Module):
"""
Overview:
The recurrent Q model for NGU(https://arxiv.org/pdf/2002.06038.pdf) policy, modified from the class DRQN in \
q_leaning.py. The implementation mentioned in the original paper is 'adapt the R2D2 agent that uses the \
dueling network architecture with an LSTM layer after a convolutional neural network'. The NGU network \
includes encoder, LSTM core(rnn) and head.
Interface:
``__init__``, ``forward``.
"""
[docs] def __init__(
self,
obs_shape: Union[int, SequenceType],
action_shape: Union[int, SequenceType],
encoder_hidden_size_list: SequenceType = [128, 128, 64],
collector_env_num: Optional[int] = 1, # TODO
dueling: bool = True,
head_hidden_size: Optional[int] = None,
head_layer_num: int = 1,
lstm_type: Optional[str] = 'normal',
activation: Optional[nn.Module] = nn.ReLU(),
norm_type: Optional[str] = None
) -> None:
"""
Overview:
Init the DRQN Model for NGU according to arguments.
Arguments:
- obs_shape (:obj:`Union[int, SequenceType]`): Observation's space, such as 8 or [4, 84, 84].
- action_shape (:obj:`Union[int, SequenceType]`): Action's space, such as 6 or [2, 3, 3].
- encoder_hidden_size_list (:obj:`SequenceType`): Collection of ``hidden_size`` to pass to ``Encoder``.
- collector_env_num (:obj:`Optional[int]`): The number of environments used to collect data simultaneously.
- dueling (:obj:`bool`): Whether choose ``DuelingHead`` (True) or ``DiscreteHead (False)``, \
default to True.
- head_hidden_size (:obj:`Optional[int]`): The ``hidden_size`` to pass to ``Head``, should match the \
last element of ``encoder_hidden_size_list``.
- head_layer_num (:obj:`int`): The number of layers in head network.
- lstm_type (:obj:`Optional[str]`): Version of rnn cell, now support ['normal', 'pytorch', 'hpc', 'gru'], \
default is 'normal'.
- 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(NGU, self).__init__()
# For compatibility: 1, (1, ), [4, H, H]
obs_shape, action_shape = squeeze(obs_shape), squeeze(action_shape)
self.action_shape = action_shape
self.collector_env_num = collector_env_num
if head_hidden_size is None:
head_hidden_size = encoder_hidden_size_list[-1]
# FC Encoder
if isinstance(obs_shape, int) or len(obs_shape) == 1:
self.encoder = FCEncoder(obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type)
# Conv Encoder
elif len(obs_shape) == 3:
self.encoder = ConvEncoder(obs_shape, encoder_hidden_size_list, activation=activation, norm_type=norm_type)
else:
raise RuntimeError(
"not support obs_shape for pre-defined encoder: {}, please customize your own DRQN".format(obs_shape)
)
# NOTE: current obs hidden_state_dim, previous action, previous extrinsic reward, beta
# TODO(pu): add prev_reward_intrinsic to network input, reward uses some kind of embedding instead of 1D value
input_size = head_hidden_size + action_shape + 1 + self.collector_env_num
# LSTM Type
self.rnn = get_lstm(lstm_type, input_size=input_size, hidden_size=head_hidden_size)
# Head Type
if dueling:
head_cls = DuelingHead
else:
head_cls = DiscreteHead
multi_head = not isinstance(action_shape, int)
if multi_head:
self.head = MultiHead(
head_cls,
head_hidden_size,
action_shape,
layer_num=head_layer_num,
activation=activation,
norm_type=norm_type
)
else:
self.head = head_cls(
head_hidden_size, action_shape, head_layer_num, activation=activation, norm_type=norm_type
)
[docs] def forward(self, inputs: Dict, inference: bool = False, saved_state_timesteps: Optional[list] = None) -> Dict:
"""
Overview:
Forward computation graph of NGU R2D2 network. Input observation, prev_action prev_reward_extrinsic \
to predict NGU Q output. Parameter updates with NGU's MLPs forward setup.
Arguments:
- inputs (:obj:`Dict`):
- obs (:obj:`torch.Tensor`): Encoded observation.
- prev_state (:obj:`list`): Previous state's tensor of size ``(B, N)``.
- inference: (:obj:'bool'): If inference is True, we unroll the one timestep transition, \
if inference is False, we unroll the sequence transitions.
- saved_state_timesteps: (:obj:'Optional[list]'): When inference is False, \
we unroll the sequence transitions, then we would save rnn hidden states at timesteps \
that are listed in list saved_state_timesteps.
Returns:
- outputs (:obj:`Dict`):
Run ``MLP`` with ``DRQN`` setups and return the result prediction dictionary.
ReturnsKeys:
- logit (:obj:`torch.Tensor`): Logit tensor with same size as input ``obs``.
- next_state (:obj:`list`): Next state's tensor of size ``(B, N)``.
Shapes:
- obs (:obj:`torch.Tensor`): :math:`(B, N=obs_space)`, where B is batch size.
- prev_state(:obj:`torch.FloatTensor list`): :math:`[(B, N)]`.
- logit (:obj:`torch.FloatTensor`): :math:`(B, N)`.
- next_state(:obj:`torch.FloatTensor list`): :math:`[(B, N)]`.
"""
x, prev_state = inputs['obs'], inputs['prev_state']
if 'prev_action' in inputs.keys():
# collect, eval mode: pass into one timestep mini-batch data (batchsize=env_num)
prev_action = inputs['prev_action']
prev_reward_extrinsic = inputs['prev_reward_extrinsic']
else:
# train mode: pass into H timesteps mini-batch data (batchsize=train_batch_size)
prev_action = torch.cat(
[torch.ones_like(inputs['action'][:, 0].unsqueeze(1)) * (-1), inputs['action'][:, :-1]], dim=1
) # (B, 1) (B, H-1) -> (B, H, self.action_shape)
prev_reward_extrinsic = torch.cat(
[torch.zeros_like(inputs['reward'][:, 0].unsqueeze(1)), inputs['reward'][:, :-1]], dim=1
) # (B, 1, nstep) (B, H-1, nstep) -> (B, H, nstep)
beta = inputs['beta'] # beta_index
if inference:
# collect, eval mode: pass into one timestep mini-batch data (batchsize=env_num)
x = self.encoder(x)
x = x.unsqueeze(0)
prev_reward_extrinsic = prev_reward_extrinsic.unsqueeze(0).unsqueeze(-1)
env_num = self.collector_env_num
beta_onehot = one_hot(beta, env_num).unsqueeze(0)
prev_action_onehot = one_hot(prev_action, self.action_shape).unsqueeze(0)
x_a_r_beta = torch.cat(
[x, prev_action_onehot, prev_reward_extrinsic, beta_onehot], dim=-1
) # shape (1, H, 1+env_num+action_dim)
x, next_state = self.rnn(x_a_r_beta.to(torch.float32), prev_state)
# TODO(pu): x, next_state = self.rnn(x, prev_state)
x = x.squeeze(0)
x = self.head(x)
x['next_state'] = next_state
return x
else:
# train mode: pass into H timesteps mini-batch data (batchsize=train_batch_size)
assert len(x.shape) in [3, 5], x.shape # (B, H, obs_dim)
x = parallel_wrapper(self.encoder)(x) # (B, H, hidden_dim)
prev_reward_extrinsic = prev_reward_extrinsic[:, :, 0].unsqueeze(-1) # (B,H,1)
env_num = self.collector_env_num
beta_onehot = one_hot(beta.view(-1), env_num).view([beta.shape[0], beta.shape[1], -1]) # (B, H, env_num)
prev_action_onehot = one_hot(prev_action.view(-1), self.action_shape).view(
[prev_action.shape[0], prev_action.shape[1], -1]
) # (B, H, action_dim)
x_a_r_beta = torch.cat(
[x, prev_action_onehot, prev_reward_extrinsic, beta_onehot], dim=-1
) # (B, H, 1+env_num+action_dim)
x = x_a_r_beta
lstm_embedding = []
# TODO(nyz) how to deal with hidden_size key-value
hidden_state_list = []
if saved_state_timesteps is not None:
saved_state = []
for t in range(x.shape[0]): # T timesteps
output, prev_state = self.rnn(x[t:t + 1], prev_state)
if saved_state_timesteps is not None and t + 1 in saved_state_timesteps:
saved_state.append(prev_state)
lstm_embedding.append(output)
# only take the hidden state h
hidden_state_list.append(torch.cat([item['h'] for item in prev_state], dim=1))
x = torch.cat(lstm_embedding, 0) # [B, H, 64]
x = parallel_wrapper(self.head)(x)
# the last timestep state including the hidden state (h) and the cell state (c)
x['next_state'] = prev_state
x['hidden_state'] = torch.cat(hidden_state_list, dim=-3)
if saved_state_timesteps is not None:
# the selected saved hidden states, including the hidden state (h) and the cell state (c)
x['saved_state'] = saved_state
return x