from typing import List, Dict, Any, Tuple, Union, Optional from collections import namedtuple import torch import copy from ding.torch_utils import Adam, to_device from ding.rl_utils import coma_data, coma_error, get_train_sample from ding.model import model_wrap from ding.utils import POLICY_REGISTRY from ding.utils.data import default_collate, default_decollate, timestep_collate from .base_policy import Policy @POLICY_REGISTRY.register('coma') class COMAPolicy(Policy): r""" Overview: Policy class of COMA algorithm. COMA is a multi model reinforcement learning algorithm Interface: _init_learn, _data_preprocess_learn, _forward_learn, _reset_learn, _state_dict_learn, _load_state_dict_learn\ _init_collect, _forward_collect, _reset_collect, _process_transition, _init_eval, _forward_eval\ _reset_eval, _get_train_sample, default_model, _monitor_vars_learn Config: == ==================== ======== ============== ======================================== ======================= ID Symbol Type Default Value Description Other(Shape) == ==================== ======== ============== ======================================== ======================= 1 ``type`` str coma | 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 True | Whether the RL algorithm is on-policy | or off-policy 4. ``priority`` bool False | Whether use priority(PER) | priority sample, | update priority 5 | ``priority_`` bool False | Whether use Importance Sampling | IS weight | ``IS_weight`` | Weight to correct biased update. 6 | ``learn.update`` int 1 | 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 7 | ``learn.target_`` float 0.001 | Target network update momentum | between[0,1] | ``update_theta`` | parameter. 8 | ``learn.discount`` float 0.99 | Reward's future discount factor, aka. | may be 1 when sparse | ``_factor`` | gamma | reward env 9 | ``learn.td_`` float 0.8 | The trade-off factor of td-lambda, | ``lambda`` | which balances 1step td and mc 10 | ``learn.value_`` float 1.0 | The loss weight of value network | policy network weight | ``weight`` | is set to 1 11 | ``learn.entropy_`` float 0.01 | The loss weight of entropy | policy network weight | ``weight`` | regularization | is set to 1 == ==================== ======== ============== ======================================== ======================= """ config = dict( # (str) RL policy register name (refer to function "POLICY_REGISTRY"). type='coma', # (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, # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, learn=dict( update_per_collect=20, batch_size=32, learning_rate=0.0005, # ============================================================== # The following configs is algorithm-specific # ============================================================== # (float) target network update weight, theta * new_w + (1 - theta) * old_w, defaults in [0, 0.1] target_update_theta=0.001, # (float) discount factor for future reward, defaults int [0, 1] discount_factor=0.99, # (float) the trade-off factor of td-lambda, which balances 1step td and mc(nstep td in practice) td_lambda=0.8, # (float) the loss weight of policy network network policy_weight=0.001, # (float) the loss weight of value network value_weight=1, # (float) the loss weight of entropy regularization entropy_weight=0.01, ), collect=dict( # (int) collect n_sample data, train model n_iteration time # n_episode=32, # (int) unroll length of a train iteration(gradient update step) unroll_len=20, ), eval=dict(), ) def default_model(self) -> Tuple[str, List[str]]: """ Overview: Return this algorithm default model setting for demonstration. Returns: - model_info (:obj:`Tuple[str, List[str]]`): model name and mode import_names .. note:: The user can define and use customized network model but must obey the same inferface definition indicated \ by import_names path. For coma, ``ding.model.coma.coma`` """ return 'coma', ['ding.model.template.coma'] def _init_learn(self) -> None: """ Overview: Init the learner model of COMAPolicy Arguments: .. note:: The _init_learn method takes the argument from the self._cfg.learn in the config file - learning_rate (:obj:`float`): The learning rate fo the optimizer - gamma (:obj:`float`): The discount factor - lambda (:obj:`float`): The lambda factor, determining the mix of bootstrapping\ vs further accumulation of multistep returns at each timestep, - value_wight(:obj:`float`): The weight of value loss in total loss - entropy_weight(:obj:`float`): The weight of entropy loss in total loss - agent_num (:obj:`int`): Since this is a multi-agent algorithm, we need to input the agent num. - batch_size (:obj:`int`): Need batch size info to init hidden_state plugins """ self._priority = self._cfg.priority self._priority_IS_weight = self._cfg.priority_IS_weight assert not self._priority, "not implemented priority in COMA" self._optimizer = Adam(self._model.parameters(), lr=self._cfg.learn.learning_rate) self._gamma = self._cfg.learn.discount_factor self._lambda = self._cfg.learn.td_lambda self._policy_weight = self._cfg.learn.policy_weight self._value_weight = self._cfg.learn.value_weight self._entropy_weight = self._cfg.learn.entropy_weight self._target_model = copy.deepcopy(self._model) self._target_model = model_wrap( self._target_model, wrapper_name='target', update_type='momentum', update_kwargs={'theta': self._cfg.learn.target_update_theta} ) self._target_model = model_wrap( self._target_model, wrapper_name='hidden_state', state_num=self._cfg.learn.batch_size, init_fn=lambda: [None for _ in range(self._cfg.model.agent_num)] ) self._learn_model = model_wrap( self._model, wrapper_name='hidden_state', state_num=self._cfg.learn.batch_size, init_fn=lambda: [None for _ in range(self._cfg.model.agent_num)] ) self._learn_model.reset() self._target_model.reset() def _data_preprocess_learn(self, data: List[Any]) -> dict: r""" Overview: Preprocess the data to fit the required data format for learning Arguments: - data (:obj:`List[Dict[str, Any]]`): the data collected from collect function, the Dict in data should contain keys including at least ['obs', 'action', 'reward'] Returns: - data (:obj:`Dict[str, Any]`): the processed data, including at least \ ['obs', 'action', 'reward', 'done', 'weight'] """ # data preprocess data = timestep_collate(data) assert set(data.keys()) > set(['obs', 'action', 'reward']) if self._cuda: data = to_device(data, self._device) data['weight'] = data.get('weight', None) data['done'] = data['done'].float() return data def _forward_learn(self, data: dict) -> Dict[str, Any]: r""" Overview: Forward and backward function of learn mode, acquire the data and calculate the loss and\ optimize learner model Arguments: - data (:obj:`Dict[str, Any]`): Dict type data, a batch of data for training, values are torch.Tensor or \ np.ndarray or dict/list combinations. Returns: - info_dict (:obj:`Dict[str, Any]`): Dict type data, a info dict indicated training result, which will be \ recorded in text log and tensorboard, values are python scalar or a list of scalars. ArgumentsKeys: - necessary: ``obs``, ``action``, ``reward``, ``done``, ``weight`` ReturnsKeys: - necessary: ``cur_lr``, ``total_loss``, ``policy_loss``, ``value_loss``, ``entropy_loss`` - cur_lr (:obj:`float`): Current learning rate - total_loss (:obj:`float`): The calculated loss - policy_loss (:obj:`float`): The policy(actor) loss of coma - value_loss (:obj:`float`): The value(critic) loss of coma - entropy_loss (:obj:`float`): The entropy loss """ data = self._data_preprocess_learn(data) # forward self._learn_model.train() self._target_model.train() self._learn_model.reset(state=data['prev_state'][0]) self._target_model.reset(state=data['prev_state'][0]) q_value = self._learn_model.forward(data, mode='compute_critic')['q_value'] with torch.no_grad(): target_q_value = self._target_model.forward(data, mode='compute_critic')['q_value'] logit = self._learn_model.forward(data, mode='compute_actor')['logit'] logit[data['obs']['action_mask'] == 0.0] = -9999999 data = coma_data(logit, data['action'], q_value, target_q_value, data['reward'], data['weight']) coma_loss = coma_error(data, self._gamma, self._lambda) total_loss = self._policy_weight * coma_loss.policy_loss + self._value_weight * coma_loss.q_value_loss - \ self._entropy_weight * coma_loss.entropy_loss # update self._optimizer.zero_grad() total_loss.backward() self._optimizer.step() # after update self._target_model.update(self._learn_model.state_dict()) return { 'cur_lr': self._optimizer.defaults['lr'], 'total_loss': total_loss.item(), 'policy_loss': coma_loss.policy_loss.item(), 'value_loss': coma_loss.q_value_loss.item(), 'entropy_loss': coma_loss.entropy_loss.item(), } def _reset_learn(self, data_id: Optional[List[int]] = None) -> None: self._learn_model.reset(data_id=data_id) 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: r""" Overview: Collect mode init moethod. Called by ``self.__init__``. Init traj and unroll length, collect model. Model has eps_greedy_sample wrapper and hidden state wrapper """ self._unroll_len = self._cfg.collect.unroll_len self._collect_model = model_wrap( self._model, wrapper_name='hidden_state', state_num=self._cfg.collect.env_num, save_prev_state=True, init_fn=lambda: [None for _ in range(self._cfg.model.agent_num)] ) self._collect_model = model_wrap(self._collect_model, wrapper_name='eps_greedy_sample') self._collect_model.reset() def _forward_collect(self, data: dict, eps: float) -> dict: r""" Overview: Collect output according to eps_greedy plugin 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]`): Dict type data, including at least inferred action according to input obs. ReturnsKeys - necessary: ``action`` """ data_id = list(data.keys()) data = default_collate(list(data.values())) if self._cuda: data = to_device(data, self._device) data = {'obs': data} self._collect_model.eval() with torch.no_grad(): output = self._collect_model.forward(data, eps=eps, data_id=data_id, mode='compute_actor') if self._cuda: output = to_device(output, 'cpu') output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} def _reset_collect(self, data_id: Optional[List[int]] = None) -> None: self._collect_model.reset(data_id=data_id) def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: r""" Overview: Generate dict type transition data from inputs. Arguments: - obs (:obj:`Any`): Env observation - model_output (:obj:`dict`): Output of collect model, including at least ['action', 'prev_state'] - timestep (:obj:`namedtuple`): Output after env step, including at least ['obs', 'reward', 'done'] \ (here 'obs' indicates obs after env step). Returns: - transition (:obj:`dict`): Dict type transition data. """ transition = { 'obs': obs, 'next_obs': timestep.obs, 'prev_state': model_output['prev_state'], 'action': model_output['action'], 'reward': timestep.reward, 'done': timestep.done, } return transition def _init_eval(self) -> None: r""" Overview: Evaluate mode init method. Called by ``self.__init__``. Init eval model with argmax strategy and hidden_state plugin. """ self._eval_model = model_wrap( self._model, wrapper_name='hidden_state', state_num=self._cfg.eval.env_num, save_prev_state=True, init_fn=lambda: [None for _ in range(self._cfg.model.agent_num)] ) self._eval_model = model_wrap(self._eval_model, wrapper_name='argmax_sample') self._eval_model.reset() def _forward_eval(self, data: dict) -> dict: r""" Overview: Forward function of eval mode, similar to ``self._forward_collect``. 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. Returns: - output (:obj:`Dict[int, Any]`): The dict of predicting action for the interaction with env. ReturnsKeys - necessary: ``action`` """ data_id = list(data.keys()) data = default_collate(list(data.values())) if self._cuda: data = to_device(data, self._device) data = {'obs': data} self._eval_model.eval() with torch.no_grad(): output = self._eval_model.forward(data, data_id=data_id, mode='compute_actor') if self._cuda: output = to_device(output, 'cpu') output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} def _reset_eval(self, data_id: Optional[List[int]] = None) -> None: self._eval_model.reset(data_id=data_id) def _get_train_sample(self, data: list) -> Union[None, List[Any]]: r""" Overview: Get the train sample from trajectory Arguments: - data (:obj:`list`): The trajectory's cache Returns: - samples (:obj:`dict`): The training samples generated """ return get_train_sample(data, self._unroll_len) def _monitor_vars_learn(self) -> List[str]: r""" Overview: Return variables' name if variables are to used in monitor. Returns: - vars (:obj:`List[str]`): Variables' name list. """ return super()._monitor_vars_learn() + ['policy_loss', 'value_loss', 'entropy_loss']