from typing import List, Dict, Any, Tuple, Union import copy import torch from ding.torch_utils import Adam, to_device from ding.rl_utils import dist_nstep_td_data, dist_nstep_td_error, get_train_sample, get_nstep_return_data from ding.model import model_wrap from ding.utils import POLICY_REGISTRY from ding.utils.data import default_collate, default_decollate from .dqn import DQNPolicy from .common_utils import default_preprocess_learn @POLICY_REGISTRY.register('c51') class C51Policy(DQNPolicy): r""" Overview: Policy class of C51 algorithm. Config: == ==================== ======== ============== ======================================== ======================= ID Symbol Type Default Value Description Other(Shape) == ==================== ======== ============== ======================================== ======================= 1 ``type`` str c51 | RL policy register name, refer to | this arg is optional, | registry ``POLICY_REGISTRY`` | a placeholder 2 ``cuda`` bool False | Whether to use cuda for network | this arg can be diff- | erent from modes 3 ``on_policy`` bool False | Whether the RL algorithm is on-policy | or off-policy 4 ``priority`` bool False | Whether use priority(PER) | priority sample, | update priority 5 ``model.v_min`` float -10 | Value of the smallest atom | in the support set. 6 ``model.v_max`` float 10 | Value of the largest atom | in the support set. 7 ``model.n_atom`` int 51 | Number of atoms in the support set | of the value distribution. 8 | ``other.eps`` float 0.95 | Start value for epsilon decay. | ``.start`` | 9 | ``other.eps`` float 0.1 | End value for epsilon decay. | ``.end`` 10 | ``discount_`` float 0.97, | Reward's future discount factor, aka. | may be 1 when sparse | ``factor`` [0.95, 0.999] | gamma | reward env 11 ``nstep`` int 1, | N-step reward discount sum for target | q_value estimation 12 | ``learn.update`` int 3 | How many updates(iterations) to train | this args can be vary | ``per_collect`` | after collector's one collection. Only | from envs. Bigger val | valid in serial training | means more off-policy == ==================== ======== ============== ======================================== ======================= """ config = dict( # (str) RL policy register name (refer to function "POLICY_REGISTRY"). type='c51', # (bool) Whether to use cuda for network. cuda=False, # (bool) Whether the RL algorithm is on-policy or off-policy. on_policy=False, # (bool) Whether use priority(priority sample, IS weight, update priority) priority=False, # (float) Reward's future discount factor, aka. gamma. discount_factor=0.97, # (int) N-step reward for target q_value estimation nstep=1, model=dict( v_min=-10, v_max=10, n_atom=51, ), learn=dict( # How many updates(iterations) to train after collector's one collection. # Bigger "update_per_collect" means bigger off-policy. # collect data -> update policy-> collect data -> ... update_per_collect=3, batch_size=64, learning_rate=0.001, # ============================================================== # The following configs are algorithm-specific # ============================================================== # (int) Frequence of target network update. target_update_freq=100, # (bool) Whether ignore done(usually for max step termination env) ignore_done=False, ), # collect_mode config collect=dict( # (int) Only one of [n_sample, n_step, n_episode] shoule be set # n_sample=8, # (int) Cut trajectories into pieces with length "unroll_len". unroll_len=1, ), eval=dict(), # other config other=dict( # Epsilon greedy with decay. eps=dict( # (str) Decay type. Support ['exp', 'linear']. type='exp', start=0.95, end=0.1, # (int) Decay length(env step) decay=10000, ), replay_buffer=dict(replay_buffer_size=10000, ) ), ) def default_model(self) -> Tuple[str, List[str]]: return 'c51dqn', ['ding.model.template.q_learning'] def _init_learn(self) -> None: r""" Overview: Learn mode init method. Called by ``self.__init__``. Init the optimizer, algorithm config, main and target models. """ self._priority = self._cfg.priority # Optimizer self._optimizer = Adam(self._model.parameters(), lr=self._cfg.learn.learning_rate) self._gamma = self._cfg.discount_factor self._nstep = self._cfg.nstep self._v_max = self._cfg.model.v_max self._v_min = self._cfg.model.v_min self._n_atom = self._cfg.model.n_atom # use wrapper instead of plugin self._target_model = copy.deepcopy(self._model) self._target_model = model_wrap( self._target_model, wrapper_name='target', update_type='assign', update_kwargs={'freq': self._cfg.learn.target_update_freq} ) self._learn_model = model_wrap(self._model, wrapper_name='argmax_sample') self._learn_model.reset() self._target_model.reset() def _forward_learn(self, data: dict) -> Dict[str, Any]: r""" Overview: Forward and backward function of learn mode. Arguments: - data (:obj:`dict`): Dict type data, including at least ['obs', 'action', 'reward', 'next_obs'] Returns: - info_dict (:obj:`Dict[str, Any]`): Including current lr and loss. """ data = default_preprocess_learn( data, use_priority=self._priority, ignore_done=self._cfg.learn.ignore_done, use_nstep=True ) if self._cuda: data = to_device(data, self._device) # ==================== # Q-learning forward # ==================== self._learn_model.train() self._target_model.train() # Current q value (main model) output = self._learn_model.forward(data['obs']) q_value = output['logit'] q_value_dist = output['distribution'] # Target q value with torch.no_grad(): target_output = self._target_model.forward(data['next_obs']) target_q_value_dist = target_output['distribution'] target_q_value = target_output['logit'] # Max q value action (main model) target_q_action = self._learn_model.forward(data['next_obs'])['action'] data_n = dist_nstep_td_data( q_value_dist, target_q_value_dist, data['action'], target_q_action, data['reward'], data['done'], data['weight'] ) value_gamma = data.get('value_gamma') loss, td_error_per_sample = dist_nstep_td_error( data_n, self._gamma, self._v_min, self._v_max, self._n_atom, nstep=self._nstep, value_gamma=value_gamma ) # ==================== # Q-learning update # ==================== self._optimizer.zero_grad() loss.backward() if self._cfg.multi_gpu: self.sync_gradients(self._learn_model) self._optimizer.step() # ============= # after update # ============= self._target_model.update(self._learn_model.state_dict()) return { 'cur_lr': self._optimizer.defaults['lr'], 'total_loss': loss.item(), 'q_value': q_value.mean().item(), 'target_q_value': target_q_value.mean().item(), 'priority': td_error_per_sample.abs().tolist(), # Only discrete action satisfying len(data['action'])==1 can return this and draw histogram on tensorboard. # '[histogram]action_distribution': data['action'], } def _monitor_vars_learn(self) -> List[str]: return ['cur_lr', 'total_loss', 'q_value', 'target_q_value'] def _state_dict_learn(self) -> Dict[str, Any]: return { 'model': self._learn_model.state_dict(), 'target_model': self._target_model.state_dict(), 'optimizer': self._optimizer.state_dict(), } def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: self._learn_model.load_state_dict(state_dict['model']) self._target_model.load_state_dict(state_dict['target_model']) self._optimizer.load_state_dict(state_dict['optimizer']) def _init_collect(self) -> None: """ Overview: Collect mode init method. Called by ``self.__init__``. Initialize necessary arguments for nstep return \ calculation and collect_model for exploration (eps_greedy_sample). """ self._unroll_len = self._cfg.collect.unroll_len self._gamma = self._cfg.discount_factor # necessary for parallel self._nstep = self._cfg.nstep # necessary for parallel self._collect_model = model_wrap(self._model, wrapper_name='eps_greedy_sample') self._collect_model.reset() def _forward_collect(self, data: Dict[int, Any], eps: float) -> Dict[int, Any]: """ Overview: Forward computation graph of collect mode(collect training data), with eps_greedy for exploration. Arguments: - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \ values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer. - eps (:obj:`float`): epsilon value for exploration, which is decayed by collected env step. Returns: - output (:obj:`Dict[int, Any]`): The dict of predicting policy_output(action) for the interaction with \ env and the constructing of transition. ArgumentsKeys: - necessary: ``obs`` ReturnsKeys - necessary: ``logit``, ``action`` """ data_id = list(data.keys()) data = default_collate(list(data.values())) if self._cuda: data = to_device(data, self._device) self._collect_model.eval() with torch.no_grad(): output = self._collect_model.forward(data, eps=eps) if self._cuda: output = to_device(output, 'cpu') output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} def _get_train_sample(self, data: list) -> Union[None, List[Any]]: """ Overview: Calculate nstep return data and transform a trajectory into many train samples. Arguments: - data (:obj:`list`): The collected data of a trajectory, which is a list that contains dict elements. Returns: - samples (:obj:`dict`): The training samples generated. """ data = get_nstep_return_data(data, self._nstep, gamma=self._gamma) return get_train_sample(data, self._unroll_len)