from typing import List, Dict, Any, Tuple, Union from collections import namedtuple import torch import copy import numpy as np from torch.distributions import Independent, Normal from ding.torch_utils import Adam, to_device, to_dtype, unsqueeze, ContrastiveLoss from ding.rl_utils import happo_data, happo_error, happo_policy_error, happo_policy_data, \ v_nstep_td_data, v_nstep_td_error, get_train_sample, gae, gae_data, happo_error_continuous, \ get_gae from ding.model import model_wrap from ding.utils import POLICY_REGISTRY, split_data_generator, RunningMeanStd from ding.utils.data import default_collate, default_decollate from .base_policy import Policy from .common_utils import default_preprocess_learn @POLICY_REGISTRY.register('happo') class HAPPOPolicy(Policy): """ Overview: Policy class of on policy version HAPPO algorithm. Paper link: https://arxiv.org/abs/2109.11251. """ config = dict( # (str) RL policy register name (refer to function "POLICY_REGISTRY"). type='happo', # (bool) Whether to use cuda for network. cuda=False, # (bool) Whether the RL algorithm is on-policy or off-policy. (Note: in practice PPO can be off-policy used) on_policy=True, # (bool) Whether to use priority(priority sample, IS weight, update priority) priority=False, # (bool) Whether to use Importance Sampling Weight to correct biased update due to priority. # If True, priority must be True. priority_IS_weight=False, # (bool) Whether to recompurete advantages in each iteration of on-policy PPO recompute_adv=True, # (str) Which kind of action space used in PPOPolicy, ['discrete', 'continuous', 'hybrid'] action_space='discrete', # (bool) Whether to use nstep return to calculate value target, otherwise, use return = adv + value nstep_return=False, # (bool) Whether to enable multi-agent training, i.e.: MAPPO multi_agent=False, # (bool) Whether to need policy data in process transition transition_with_policy_data=True, learn=dict( epoch_per_collect=10, batch_size=64, learning_rate=3e-4, # ============================================================== # The following configs is algorithm-specific # ============================================================== # (float) The loss weight of value network, policy network weight is set to 1 value_weight=0.5, # (float) The loss weight of entropy regularization, policy network weight is set to 1 entropy_weight=0.0, # (float) PPO clip ratio, defaults to 0.2 clip_ratio=0.2, # (bool) Whether to use advantage norm in a whole training batch adv_norm=True, value_norm=True, ppo_param_init=True, grad_clip_type='clip_norm', grad_clip_value=0.5, ignore_done=False, ), collect=dict( # (int) Only one of [n_sample, n_episode] shoule be set # n_sample=64, # (int) Cut trajectories into pieces with length "unroll_len". unroll_len=1, # ============================================================== # The following configs is algorithm-specific # ============================================================== # (float) Reward's future discount factor, aka. gamma. discount_factor=0.99, # (float) GAE lambda factor for the balance of bias and variance(1-step td and mc) gae_lambda=0.95, ), eval=dict(), ) def _init_learn(self) -> None: """ Overview: Initialize the learn mode of policy, including related attributes and modules. For HAPPO, it mainly \ contains optimizer, algorithm-specific arguments such as loss weight, clip_ratio and recompute_adv. This \ method also executes some special network initializations and prepares running mean/std monitor for value. This method will be called in ``__init__`` method if ``learn`` field is in ``enable_field``. .. note:: For the member variables that need to be saved and loaded, please refer to the ``_state_dict_learn`` \ and ``_load_state_dict_learn`` methods. .. note:: For the member variables that need to be monitored, please refer to the ``_monitor_vars_learn`` method. .. note:: If you want to set some spacial member variables in ``_init_learn`` method, you'd better name them \ with prefix ``_learn_`` to avoid conflict with other modes, such as ``self._learn_attr1``. """ self._priority = self._cfg.priority self._priority_IS_weight = self._cfg.priority_IS_weight assert not self._priority and not self._priority_IS_weight, "Priority is not implemented in PPO" assert self._cfg.action_space in ["continuous", "discrete"] self._action_space = self._cfg.action_space if self._cfg.learn.ppo_param_init: for n, m in self._model.named_modules(): if isinstance(m, torch.nn.Linear): torch.nn.init.orthogonal_(m.weight) torch.nn.init.zeros_(m.bias) if self._action_space in ['continuous']: # init log sigma for agent_id in range(self._cfg.agent_num): # if hasattr(self._model.agent_models[agent_id].actor_head, 'log_sigma_param'): # torch.nn.init.constant_(self._model.agent_models[agent_id].actor_head.log_sigma_param, 1) # The above initialization step has been changed to reparameterizationHead. for m in list(self._model.agent_models[agent_id].critic.modules()) + \ list(self._model.agent_models[agent_id].actor.modules()): if isinstance(m, torch.nn.Linear): # orthogonal initialization torch.nn.init.orthogonal_(m.weight, gain=np.sqrt(2)) torch.nn.init.zeros_(m.bias) # do last policy layer scaling, this will make initial actions have (close to) # 0 mean and std, and will help boost performances, # see https://arxiv.org/abs/2006.05990, Fig.24 for details for m in self._model.agent_models[agent_id].actor.modules(): if isinstance(m, torch.nn.Linear): torch.nn.init.zeros_(m.bias) m.weight.data.copy_(0.01 * m.weight.data) # Add the actor/critic parameters of each HAVACAgent in HAVAC to the parameter list of actor/critic_optimizer actor_params = [] critic_params = [] for agent_idx in range(self._model.agent_num): actor_params.append({'params': self._model.agent_models[agent_idx].actor.parameters()}) critic_params.append({'params': self._model.agent_models[agent_idx].critic.parameters()}) self._actor_optimizer = Adam( actor_params, lr=self._cfg.learn.learning_rate, grad_clip_type=self._cfg.learn.grad_clip_type, clip_value=self._cfg.learn.grad_clip_value, # eps = 1e-5, ) self._critic_optimizer = Adam( critic_params, lr=self._cfg.learn.critic_learning_rate, grad_clip_type=self._cfg.learn.grad_clip_type, clip_value=self._cfg.learn.grad_clip_value, # eps = 1e-5, ) self._learn_model = model_wrap(self._model, wrapper_name='base') # 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)] # ) # Algorithm config self._value_weight = self._cfg.learn.value_weight self._entropy_weight = self._cfg.learn.entropy_weight self._clip_ratio = self._cfg.learn.clip_ratio self._adv_norm = self._cfg.learn.adv_norm self._value_norm = self._cfg.learn.value_norm if self._value_norm: self._running_mean_std = RunningMeanStd(epsilon=1e-4, device=self._device) self._gamma = self._cfg.collect.discount_factor self._gae_lambda = self._cfg.collect.gae_lambda self._recompute_adv = self._cfg.recompute_adv # Main model self._learn_model.reset() def prepocess_data_agent(self, data: Dict[str, Any]): """ Overview: Preprocess data for agent dim. This function is used in learn mode. \ It will be called recursively to process nested dict data. \ It will transpose the data with shape (B, agent_num, ...) to (agent_num, B, ...). \ Arguments: - data (:obj:`dict`): Dict type data, where each element is the data of an agent of dict type. Returns: - ret (:obj:`dict`): Dict type data, where each element is the data of an agent of dict type. """ ret = {} for key, value in data.items(): if isinstance(value, dict): ret[key] = self.prepocess_data_agent(value) elif isinstance(value, torch.Tensor) and len(value.shape) > 1: ret[key] = value.transpose(0, 1) else: ret[key] = value return ret def _forward_learn(self, data: Dict[str, Any]) -> Dict[str, Any]: """ Overview: Forward and backward function of learn mode. Arguments: - data (:obj:`dict`): List type data, where each element is the data of an agent of dict type. Returns: - info_dict (:obj:`Dict[str, Any]`): Including current lr, total_loss, policy_loss, value_loss, entropy_loss, \ adv_abs_max, approx_kl, clipfrac Overview: Policy forward function of learn mode (training policy and updating parameters). Forward means \ that the policy inputs some training batch data from the replay buffer and then returns the output \ result, including various training information such as loss, clipfrac, approx_kl. Arguments: - data (:obj:`List[Dict[int, Any]]`): The input data used for policy forward, including the latest \ collected training samples for on-policy algorithms like HAPPO. For each element in list, the key of \ dict is the name of data items and the value is the corresponding data. Usually, the value is \ torch.Tensor or np.ndarray or there dict/list combinations. In the ``_forward_learn`` method, data \ often need to first be stacked in the batch dimension by some utility functions such as \ ``default_preprocess_learn``. \ For HAPPO, each element in list is a dict containing at least the following keys: ``obs``, \ ``action``, ``reward``, ``logit``, ``value``, ``done``. Sometimes, it also contains other keys \ such as ``weight``. Returns: - return_infos (:obj:`List[Dict[str, Any]]`): The information list that indicated training result, each \ training iteration contains append a information dict into the final list. The list will be precessed \ and recorded in text log and tensorboard. The value of the dict must be python scalar or a list of \ scalars. For the detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. .. tip:: The training procedure of HAPPO is three for loops. The outermost loop trains each agent separately. \ The middle loop trains all the collected training samples with ``epoch_per_collect`` epochs. The inner \ loop splits all the data into different mini-batch with the length of ``batch_size``. .. note:: The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ For the data type that not supported, the main reason is that the corresponding model does not support it. \ You can implement you own model rather than use the default model. For more information, please raise an \ issue in GitHub repo and we will continue to follow up. .. note:: For more detailed examples, please refer to our unittest for HAPPOPolicy: ``ding.policy.tests.test_happo``. """ data = default_preprocess_learn(data, ignore_done=self._cfg.learn.ignore_done, use_nstep=False) all_data_len = data['obs']['agent_state'].shape[0] # fator is the ratio of the old and new strategies of the first m-1 agents, initialized to 1. # Each transition has its own factor. ref: http://arxiv.org/abs/2109.11251 factor = torch.ones(all_data_len, 1) # (B, 1) if self._cuda: data = to_device(data, self._device) factor = to_device(factor, self._device) # process agent dim data = self.prepocess_data_agent(data) # ==================== # PPO forward # ==================== return_infos = [] self._learn_model.train() for agent_id in range(self._cfg.agent_num): agent_data = {} for key, value in data.items(): if value is not None: if type(value) is dict: agent_data[key] = {k: v[agent_id] for k, v in value.items()} # not feasible for rnn elif len(value.shape) > 1: agent_data[key] = data[key][agent_id] else: agent_data[key] = data[key] else: agent_data[key] = data[key] # update factor agent_data['factor'] = factor # calculate old_logits of all data in buffer for later factor inputs = { 'obs': agent_data['obs'], # 'actor_prev_state': agent_data['actor_prev_state'], # 'critic_prev_state': agent_data['critic_prev_state'], } old_logits = self._learn_model.forward(agent_id, inputs, mode='compute_actor')['logit'] for epoch in range(self._cfg.learn.epoch_per_collect): if self._recompute_adv: # calculate new value using the new updated value network with torch.no_grad(): inputs['obs'] = agent_data['obs'] # value = self._learn_model.forward(agent_id, agent_data['obs'], mode='compute_critic')['value'] value = self._learn_model.forward(agent_id, inputs, mode='compute_critic')['value'] inputs['obs'] = agent_data['next_obs'] next_value = self._learn_model.forward(agent_id, inputs, mode='compute_critic')['value'] if self._value_norm: value *= self._running_mean_std.std next_value *= self._running_mean_std.std traj_flag = agent_data.get('traj_flag', None) # traj_flag indicates termination of trajectory compute_adv_data = gae_data( value, next_value, agent_data['reward'], agent_data['done'], traj_flag ) agent_data['adv'] = gae(compute_adv_data, self._gamma, self._gae_lambda) unnormalized_returns = value + agent_data['adv'] if self._value_norm: agent_data['value'] = value / self._running_mean_std.std agent_data['return'] = unnormalized_returns / self._running_mean_std.std self._running_mean_std.update(unnormalized_returns.cpu().numpy()) else: agent_data['value'] = value agent_data['return'] = unnormalized_returns else: # don't recompute adv if self._value_norm: unnormalized_return = agent_data['adv'] + agent_data['value'] * self._running_mean_std.std agent_data['return'] = unnormalized_return / self._running_mean_std.std self._running_mean_std.update(unnormalized_return.cpu().numpy()) else: agent_data['return'] = agent_data['adv'] + agent_data['value'] for batch in split_data_generator(agent_data, self._cfg.learn.batch_size, shuffle=True): inputs = { 'obs': batch['obs'], # 'actor_prev_state': batch['actor_prev_state'], # 'critic_prev_state': batch['critic_prev_state'], } output = self._learn_model.forward(agent_id, inputs, mode='compute_actor_critic') adv = batch['adv'] if self._adv_norm: # Normalize advantage in a train_batch adv = (adv - adv.mean()) / (adv.std() + 1e-8) # Calculate happo error if self._action_space == 'continuous': happo_batch = happo_data( output['logit'], batch['logit'], batch['action'], output['value'], batch['value'], adv, batch['return'], batch['weight'], batch['factor'] ) happo_loss, happo_info = happo_error_continuous(happo_batch, self._clip_ratio) elif self._action_space == 'discrete': happo_batch = happo_data( output['logit'], batch['logit'], batch['action'], output['value'], batch['value'], adv, batch['return'], batch['weight'], batch['factor'] ) happo_loss, happo_info = happo_error(happo_batch, self._clip_ratio) wv, we = self._value_weight, self._entropy_weight total_loss = happo_loss.policy_loss + wv * happo_loss.value_loss - we * happo_loss.entropy_loss # actor update # critic update self._actor_optimizer.zero_grad() self._critic_optimizer.zero_grad() total_loss.backward() self._actor_optimizer.step() self._critic_optimizer.step() return_info = { 'agent{}_cur_lr'.format(agent_id): self._actor_optimizer.defaults['lr'], 'agent{}_total_loss'.format(agent_id): total_loss.item(), 'agent{}_policy_loss'.format(agent_id): happo_loss.policy_loss.item(), 'agent{}_value_loss'.format(agent_id): happo_loss.value_loss.item(), 'agent{}_entropy_loss'.format(agent_id): happo_loss.entropy_loss.item(), 'agent{}_adv_max'.format(agent_id): adv.max().item(), 'agent{}_adv_mean'.format(agent_id): adv.mean().item(), 'agent{}_value_mean'.format(agent_id): output['value'].mean().item(), 'agent{}_value_max'.format(agent_id): output['value'].max().item(), 'agent{}_approx_kl'.format(agent_id): happo_info.approx_kl, 'agent{}_clipfrac'.format(agent_id): happo_info.clipfrac, } if self._action_space == 'continuous': return_info.update( { 'agent{}_act'.format(agent_id): batch['action'].float().mean().item(), 'agent{}_mu_mean'.format(agent_id): output['logit']['mu'].mean().item(), 'agent{}_sigma_mean'.format(agent_id): output['logit']['sigma'].mean().item(), } ) return_infos.append(return_info) # calculate the factor inputs = { 'obs': agent_data['obs'], # 'actor_prev_state': agent_data['actor_prev_state'], } new_logits = self._learn_model.forward(agent_id, inputs, mode='compute_actor')['logit'] if self._cfg.action_space == 'discrete': dist_new = torch.distributions.categorical.Categorical(logits=new_logits) dist_old = torch.distributions.categorical.Categorical(logits=old_logits) elif self._cfg.action_space == 'continuous': dist_new = Normal(new_logits['mu'], new_logits['sigma']) dist_old = Normal(old_logits['mu'], old_logits['sigma']) logp_new = dist_new.log_prob(agent_data['action']) logp_old = dist_old.log_prob(agent_data['action']) if len(logp_new.shape) > 1: # for logp with shape(B, action_shape), we need to calculate the product of all action dimensions. factor = factor * torch.prod( torch.exp(logp_new - logp_old), dim=-1 ).reshape(all_data_len, 1).detach() # attention the shape else: # for logp with shape(B, ), directly calculate factor factor = factor * torch.exp(logp_new - logp_old).reshape(all_data_len, 1).detach() return return_infos def _state_dict_learn(self) -> Dict[str, Any]: """ Overview: Return the state_dict of learn mode optimizer and model. Returns: - state_dict (:obj:`Dict[str, Any]`): The dict of current policy learn mode. It contains the \ state_dict of current policy network and optimizer. """ return { 'model': self._learn_model.state_dict(), 'actor_optimizer': self._actor_optimizer.state_dict(), 'critic_optimizer': self._critic_optimizer.state_dict(), } def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: """ Overview: Load the state_dict of learn mode optimizer and model. Arguments: - state_dict (:obj:`Dict[str, Any]`): The dict of policy learn mode. It contains the state_dict \ of current policy network and optimizer. """ self._learn_model.load_state_dict(state_dict['model']) self._actor_optimizer.load_state_dict(state_dict['actor_optimizer']) self._critic_optimizer.load_state_dict(state_dict['critic_optimizer']) def _init_collect(self) -> None: """ Overview: Initialize the collect mode of policy, including related attributes and modules. For HAPPO, it contains \ the collect_model to balance the exploration and exploitation (e.g. the multinomial sample mechanism in \ discrete action space), and other algorithm-specific arguments such as unroll_len and gae_lambda. This method will be called in ``__init__`` method if ``collect`` field is in ``enable_field``. .. note:: If you want to set some spacial member variables in ``_init_collect`` method, you'd better name them \ with prefix ``_collect_`` to avoid conflict with other modes, such as ``self._collect_attr1``. .. tip:: Some variables need to initialize independently in different modes, such as gamma and gae_lambda in PPO. \ This design is for the convenience of parallel execution of different policy modes. """ self._unroll_len = self._cfg.collect.unroll_len assert self._cfg.action_space in ["continuous", "discrete"] self._action_space = self._cfg.action_space if self._action_space == 'continuous': self._collect_model = model_wrap(self._model, wrapper_name='reparam_sample') elif self._action_space == 'discrete': self._collect_model = model_wrap(self._model, wrapper_name='multinomial_sample') self._collect_model.reset() self._gamma = self._cfg.collect.discount_factor self._gae_lambda = self._cfg.collect.gae_lambda self._recompute_adv = self._cfg.recompute_adv def _forward_collect(self, data: Dict[int, Any]) -> dict: """ Overview: Policy forward function of collect mode (collecting training data by interacting with envs). Forward means \ that the policy gets some necessary data (mainly observation) from the envs and then returns the output \ data, such as the action to interact with the envs. Arguments: - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ key of the dict is environment id and the value is the corresponding data of the env. Returns: - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action and \ other necessary data (action logit and value) for learn mode defined in ``self._process_transition`` \ method. The key of the dict is the same as the input data, i.e. environment id. .. tip:: If you want to add more tricks on this policy, like temperature factor in multinomial sample, you can pass \ related data as extra keyword arguments of this method. .. note:: The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ For the data type that not supported, the main reason is that the corresponding model does not support it. \ You can implement you own model rather than use the default model. For more information, please raise an \ issue in GitHub repo and we will continue to follow up. .. note:: For more detailed examples, please refer to our unittest for HAPPOPolicy: ``ding.policy.tests.test_happo``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) if self._cuda: data = to_device(data, self._device) data = {k: v.transpose(0, 1) for k, v in data.items()} # not feasible for rnn self._collect_model.eval() with torch.no_grad(): outputs = [] for agent_id in range(self._cfg.agent_num): # output = self._collect_model.forward(agent_id, data, mode='compute_actor_critic') single_agent_obs = {k: v[agent_id] for k, v in data.items()} input = { 'obs': single_agent_obs, } output = self._collect_model.forward(agent_id, input, mode='compute_actor_critic') outputs.append(output) # transfer data from (M, B, N)->(B, M, N) result = {} for key in outputs[0].keys(): if isinstance(outputs[0][key], dict): subkeys = outputs[0][key].keys() stacked_subvalues = {} for subkey in subkeys: stacked_subvalues[subkey] = \ torch.stack([output[key][subkey] for output in outputs], dim=0).transpose(0, 1) result[key] = stacked_subvalues else: # If Value is tensor, stack it directly if isinstance(outputs[0][key], torch.Tensor): result[key] = torch.stack([output[key] for output in outputs], dim=0).transpose(0, 1) else: # If it is not tensor, assume that it is a non-stackable data type \ # (such as int, float, etc.), and directly retain the original value result[key] = [output[key] for output in outputs] output = result if self._cuda: output = to_device(output, 'cpu') output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict: """ Overview: Process and pack one timestep transition data into a dict, which can be directly used for training and \ saved in replay buffer. For HAPPO, it contains obs, next_obs, action, reward, done, logit, value. Arguments: - obs (:obj:`torch.Tensor`): The env observation of current timestep. - policy_output (:obj:`Dict[str, torch.Tensor]`): The output of the policy network with the observation \ as input. For PPO, it contains the state value, action and the logit of the action. - timestep (:obj:`namedtuple`): The execution result namedtuple returned by the environment step method, \ except all the elements have been transformed into tensor data. Usually, it contains the next obs, \ reward, done, info, etc. Returns: - transition (:obj:`Dict[str, torch.Tensor]`): The processed transition data of the current timestep. .. note:: ``next_obs`` is used to calculate nstep return when necessary, so we place in into transition by default. \ You can delete this field to save memory occupancy if you do not need nstep return. """ transition = { 'obs': obs, 'next_obs': timestep.obs, 'action': model_output['action'], 'logit': model_output['logit'], 'value': model_output['value'], 'reward': timestep.reward, 'done': timestep.done, } return transition def _get_train_sample(self, data: list) -> Union[None, List[Any]]: """ Overview: For a given trajectory (transitions, a list of transition) data, process it into a list of sample that \ can be used for training directly. In HAPPO, a train sample is a processed transition with new computed \ ``traj_flag`` and ``adv`` field. This method is usually used in collectors to execute necessary \ RL data preprocessing before training, which can help learner amortize revelant time consumption. \ In addition, you can also implement this method as an identity function and do the data processing \ in ``self._forward_learn`` method. Arguments: - transitions (:obj:`List[Dict[str, Any]`): The trajectory data (a list of transition), each element is \ the same format as the return value of ``self._process_transition`` method. Returns: - samples (:obj:`List[Dict[str, Any]]`): The processed train samples, each element is the similar format \ as input transitions, but may contain more data for training, such as GAE advantage. """ data = to_device(data, self._device) for transition in data: transition['traj_flag'] = copy.deepcopy(transition['done']) data[-1]['traj_flag'] = True if self._cfg.learn.ignore_done: data[-1]['done'] = False if data[-1]['done']: last_value = torch.zeros_like(data[-1]['value']) else: with torch.no_grad(): last_values = [] for agent_id in range(self._cfg.agent_num): inputs = {'obs': {k: unsqueeze(v[agent_id], 0) for k, v in data[-1]['next_obs'].items()}} last_value = self._collect_model.forward(agent_id, inputs, mode='compute_actor_critic')['value'] last_values.append(last_value) last_value = torch.cat(last_values) if len(last_value.shape) == 2: # multi_agent case: last_value = last_value.squeeze(0) if self._value_norm: last_value *= self._running_mean_std.std for i in range(len(data)): data[i]['value'] *= self._running_mean_std.std data = get_gae( data, to_device(last_value, self._device), gamma=self._gamma, gae_lambda=self._gae_lambda, cuda=False, ) if self._value_norm: for i in range(len(data)): data[i]['value'] /= self._running_mean_std.std # remove next_obs for save memory when not recompute adv if not self._recompute_adv: for i in range(len(data)): data[i].pop('next_obs') return get_train_sample(data, self._unroll_len) def _init_eval(self) -> None: """ Overview: Initialize the eval mode of policy, including related attributes and modules. For PPO, it contains the \ eval model to select optimial action (e.g. greedily select action with argmax mechanism in discrete action). This method will be called in ``__init__`` method if ``eval`` field is in ``enable_field``. .. note:: If you want to set some spacial member variables in ``_init_eval`` method, you'd better name them \ with prefix ``_eval_`` to avoid conflict with other modes, such as ``self._eval_attr1``. """ assert self._cfg.action_space in ["continuous", "discrete"] self._action_space = self._cfg.action_space if self._action_space == 'continuous': self._eval_model = model_wrap(self._model, wrapper_name='deterministic_sample') elif self._action_space == 'discrete': self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample') self._eval_model.reset() def _forward_eval(self, data: dict) -> dict: """ Overview: Policy forward function of eval mode (evaluation policy performance by interacting with envs). Forward \ means that the policy gets some necessary data (mainly observation) from the envs and then returns the \ action to interact with the envs. ``_forward_eval`` in HAPPO often uses deterministic sample method to \ get actions while ``_forward_collect`` usually uses stochastic sample method for balance exploration and \ exploitation. Arguments: - data (:obj:`Dict[int, Any]`): The input data used for policy forward, including at least the obs. The \ key of the dict is environment id and the value is the corresponding data of the env. Returns: - output (:obj:`Dict[int, Any]`): The output data of policy forward, including at least the action. The \ key of the dict is the same as the input data, i.e. environment id. .. note:: The input value can be torch.Tensor or dict/list combinations and current policy supports all of them. \ For the data type that not supported, the main reason is that the corresponding model does not support it. \ You can implement you own model rather than use the default model. For more information, please raise an \ issue in GitHub repo and we will continue to follow up. .. note:: For more detailed examples, please refer to our unittest for HAPPOPolicy: ``ding.policy.tests.test_happo``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) if self._cuda: data = to_device(data, self._device) # transfer data from (B, M, N)->(M, B, N) data = {k: v.transpose(0, 1) for k, v in data.items()} # not feasible for rnn self._eval_model.eval() with torch.no_grad(): outputs = [] for agent_id in range(self._cfg.agent_num): single_agent_obs = {k: v[agent_id] for k, v in data.items()} input = { 'obs': single_agent_obs, } output = self._eval_model.forward(agent_id, input, mode='compute_actor') outputs.append(output) output = self.revert_agent_data(outputs) if self._cuda: output = to_device(output, 'cpu') output = default_decollate(output) return {i: d for i, d in zip(data_id, output)} def default_model(self) -> Tuple[str, List[str]]: """ Overview: Return this algorithm default neural network model setting for demonstration. ``__init__`` method will \ automatically call this method to get the default model setting and create model. Returns: - model_info (:obj:`Tuple[str, List[str]]`): The registered model name and model's 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 example about HAPPO, its registered name is ``happo`` and the import_names is \ ``ding.model.template.havac``. """ return 'havac', ['ding.model.template.havac'] def _monitor_vars_learn(self) -> List[str]: """ Overview: Return the necessary keys for logging the return dict of ``self._forward_learn``. The logger module, such \ as text logger, tensorboard logger, will use these keys to save the corresponding data. Returns: - necessary_keys (:obj:`List[str]`): The list of the necessary keys to be logged. """ variables = super()._monitor_vars_learn() + [ 'policy_loss', 'value_loss', 'entropy_loss', 'adv_max', 'adv_mean', 'approx_kl', 'clipfrac', 'value_max', 'value_mean', ] if self._action_space == 'continuous': variables += ['mu_mean', 'sigma_mean', 'sigma_grad', 'act'] prefixes = [f'agent{i}_' for i in range(self._cfg.agent_num)] variables = [prefix + var for prefix in prefixes for var in variables] return variables def revert_agent_data(self, data: list): """ Overview: Revert the data of each agent to the original data format. Arguments: - data (:obj:`list`): List type data, where each element is the data of an agent of dict type. Returns: - ret (:obj:`dict`): Dict type data, where each element is the data of an agent of dict type. """ ret = {} # Traverse all keys of the first output for key in data[0].keys(): if isinstance(data[0][key], torch.Tensor): # If the value corresponding to the current key is tensor, stack N tensors stacked_tensor = torch.stack([output[key] for output in data], dim=0) ret[key] = stacked_tensor.transpose(0, 1) elif isinstance(data[0][key], dict): # If the value corresponding to the current key is a dictionary, recursively \ # call the function to process the contents inside the dictionary. ret[key] = self.revert_agent_data([output[key] for output in data]) return ret