from typing import List, Dict, Any, Tuple, Union from collections import namedtuple import copy import torch from torch.utils.data import Dataset, DataLoader from ding.utils import POLICY_REGISTRY, split_data_generator, RunningMeanStd from ding.utils.data import default_collate, default_decollate from ding.torch_utils import Adam, to_device from ding.rl_utils import get_gae_with_default_last_value, get_train_sample, gae, gae_data, get_gae, \ ppo_policy_data, ppo_policy_error, ppo_value_data, ppo_value_error, ppg_data, ppg_joint_error from ding.model import model_wrap from .base_policy import Policy class ExperienceDataset(Dataset): """ Overview: A dataset class for storing and accessing experience data. Interface: ``__init__``, ``__len__``, ``__getitem__``. """ def __init__(self, data): """ Arguments: - data (:obj:`dict`): A dictionary containing the experience data, where the keys represent the data types \ and the values are the corresponding data arrays. """ super().__init__() self.data = data def __len__(self): return list(self.data.values())[0].shape[0] def __getitem__(self, ind): data = {} for key in self.data.keys(): data[key] = self.data[key][ind] return data def create_shuffled_dataloader(data, batch_size): ds = ExperienceDataset(data) return DataLoader(ds, batch_size=batch_size, shuffle=True) @POLICY_REGISTRY.register('ppg') class PPGPolicy(Policy): """ Overview: Policy class of PPG algorithm. PPG is a policy gradient algorithm with auxiliary phase training. \ The auxiliary phase training is proposed to distill the value into the policy network, \ while making sure the policy network does not change the action predictions (kl div loss). \ Paper link: https://arxiv.org/abs/2009.04416. Interface: ``_init_learn``, ``_data_preprocess_learn``, ``_forward_learn``, ``_state_dict_learn``, \ ``_load_state_dict_learn``, ``_init_collect``, ``_forward_collect``, ``_process_transition``, \ ``_get_train_sample``, ``_get_batch_size``, ``_init_eval``, ``_forward_eval``, ``default_model``, \ ``_monitor_vars_learn``, ``learn_aux``. Config: == ==================== ======== ============== ======================================== ======================= ID Symbol Type Default Value Description Other(Shape) == ==================== ======== ============== ======================================== ======================= 1 ``type`` str ppg | 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 5 | 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.value_`` float 1.0 | The loss weight of value network | policy network weight | ``weight`` | is set to 1 8 | ``learn.entropy_`` float 0.01 | The loss weight of entropy | policy network weight | ``weight`` | regularization | is set to 1 9 | ``learn.clip_`` float 0.2 | PPO clip ratio | ``ratio`` 10 | ``learn.adv_`` bool False | Whether to use advantage norm in | ``norm`` | a whole training batch 11 | ``learn.aux_`` int 5 | The frequency(normal update times) | ``freq`` | of auxiliary phase training 12 | ``learn.aux_`` int 6 | The training epochs of auxiliary | ``train_epoch`` | phase 13 | ``learn.aux_`` int 1 | The loss weight of behavioral_cloning | ``bc_weight`` | in auxiliary phase 14 | ``collect.dis`` float 0.99 | Reward's future discount factor, aka. | may be 1 when sparse | ``count_factor`` | gamma | reward env 15 | ``collect.gae_`` float 0.95 | GAE lambda factor for the balance | ``lambda`` | of bias and variance(1-step td and mc) == ==================== ======== ============== ======================================== ======================= """ config = dict( # (str) RL policy register name (refer to function "POLICY_REGISTRY"). type='ppg', # (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, priority=False, # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, learn=dict( actor_epoch_per_collect=1, critic_epoch_per_collect=1, batch_size=64, learning_rate=0.001, # ============================================================== # 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.01, # (float) PPO clip ratio, defaults to 0.2 clip_ratio=0.2, value_norm=False, # (bool) Whether to use advantage norm in a whole training batch adv_norm=False, # (int) The frequency(normal update times) of auxiliary phase training aux_freq=8, # (int) The training epochs of auxiliary phase aux_train_epoch=6, # (int) The loss weight of behavioral_cloning in auxiliary phase aux_bc_weight=1, grad_clip_type='clip_norm', grad_clip_value=10, ignore_done=False, ), collect=dict( # n_sample=64, 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 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. """ return 'ppg', ['ding.model.template.ppg'] def _init_learn(self) -> None: """ Overview: Initialize the learn mode of policy, including related attributes and modules. For PPG, it mainly \ contains optimizer, algorithm-specific arguments such as aux_bc_weight and aux_train_epoch. 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``. """ # Optimizer self._optimizer_ac = Adam(self._model.actor_critic.parameters(), lr=self._cfg.learn.learning_rate) self._optimizer_aux_critic = Adam(self._model.aux_critic.parameters(), lr=self._cfg.learn.learning_rate) self._learn_model = model_wrap(self._model, wrapper_name='base') # Algorithm config 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 PPG" self._value_weight = self._cfg.learn.value_weight self._entropy_weight = self._cfg.learn.entropy_weight self._value_norm = self._cfg.learn.value_norm if self._value_norm: self._running_mean_std = RunningMeanStd(epsilon=1e-4, device=self._device) self._clip_ratio = self._cfg.learn.clip_ratio self._adv_norm = self._cfg.learn.adv_norm # Main model self._learn_model.reset() # Auxiliary memories self._aux_train_epoch = self._cfg.learn.aux_train_epoch self._train_iteration = 0 self._aux_memories = [] self._aux_bc_weight = self._cfg.learn.aux_bc_weight def _data_preprocess_learn(self, data: List[Any]) -> dict: """ Overview: Preprocess the data to fit the required data format for learning, including \ collate(stack data into batch), ignore done(in some fake terminate env),\ prepare loss weight per training sample, and cpu tensor to cuda. Arguments: - data (:obj:`List[Dict[str, Any]]`): The data collected from collect function. Returns: - data (:obj:`Dict[str, Any]`): The processed data, including at least ['done', 'weight']. """ # data preprocess data = default_collate(data) ignore_done = self._cfg.learn.ignore_done if ignore_done: data['done'] = None else: data['done'] = data['done'].float() data['weight'] = None if self._cuda: data = to_device(data, self._device) return data def _forward_learn(self, data: dict) -> Dict[str, Any]: """ Overview: Forward and backward function of learn mode. Arguments: - data (:obj:`Dict[str, Any]`): Input data used for policy forward, including the \ collected training samples from replay buffer. For each element in dict, the key of the \ 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 PPG, 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: - 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. \ For the detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` 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 PPGPolicy: ``ding.policy.tests.test_ppgs``. """ data = self._data_preprocess_learn(data) # ==================== # PPG forward # ==================== self._learn_model.train() return_infos = [] if self._value_norm: unnormalized_return = data['adv'] + data['value'] * self._running_mean_std.std data['return'] = unnormalized_return / self._running_mean_std.std self._running_mean_std.update(unnormalized_return.cpu().numpy()) else: data['return'] = data['adv'] + data['value'] for epoch in range(self._cfg.learn.actor_epoch_per_collect): for policy_data in split_data_generator(data, self._cfg.learn.batch_size, shuffle=True): policy_adv = policy_data['adv'] if self._adv_norm: # Normalize advantage in a total train_batch policy_adv = (policy_adv - policy_adv.mean()) / (policy_adv.std() + 1e-8) # Policy Phase(Policy) policy_output = self._learn_model.forward(policy_data['obs'], mode='compute_actor') policy_error_data = ppo_policy_data( policy_output['logit'], policy_data['logit'], policy_data['action'], policy_adv, policy_data['weight'] ) ppo_policy_loss, ppo_info = ppo_policy_error(policy_error_data, self._clip_ratio) policy_loss = ppo_policy_loss.policy_loss - self._entropy_weight * ppo_policy_loss.entropy_loss self._optimizer_ac.zero_grad() policy_loss.backward() self._optimizer_ac.step() for epoch in range(self._cfg.learn.critic_epoch_per_collect): for value_data in split_data_generator(data, self._cfg.learn.batch_size, shuffle=True): value_adv = value_data['adv'] return_ = value_data['return'] if self._adv_norm: # Normalize advantage in a total train_batch value_adv = (value_adv - value_adv.mean()) / (value_adv.std() + 1e-8) # Policy Phase(Value) value_output = self._learn_model.forward(value_data['obs'], mode='compute_critic') value_error_data = ppo_value_data( value_output['value'], value_data['value'], return_, value_data['weight'] ) value_loss = self._value_weight * ppo_value_error(value_error_data, self._clip_ratio) self._optimizer_aux_critic.zero_grad() value_loss.backward() self._optimizer_aux_critic.step() data['return_'] = data['return'] self._aux_memories.append(copy.deepcopy(data)) self._train_iteration += 1 # ==================== # PPG update # use aux loss after iterations and reset aux_memories # ==================== # Auxiliary Phase # record data for auxiliary head if self._train_iteration % self._cfg.learn.aux_freq == 0: aux_loss, bc_loss, aux_value_loss = self.learn_aux() return { 'policy_cur_lr': self._optimizer_ac.defaults['lr'], 'value_cur_lr': self._optimizer_aux_critic.defaults['lr'], 'policy_loss': ppo_policy_loss.policy_loss.item(), 'value_loss': value_loss.item(), 'entropy_loss': ppo_policy_loss.entropy_loss.item(), 'policy_adv_abs_max': policy_adv.abs().max().item(), 'approx_kl': ppo_info.approx_kl, 'clipfrac': ppo_info.clipfrac, 'aux_value_loss': aux_value_loss, 'auxiliary_loss': aux_loss, 'behavioral_cloning_loss': bc_loss, } else: return { 'policy_cur_lr': self._optimizer_ac.defaults['lr'], 'value_cur_lr': self._optimizer_aux_critic.defaults['lr'], 'policy_loss': ppo_policy_loss.policy_loss.item(), 'value_loss': value_loss.item(), 'entropy_loss': ppo_policy_loss.entropy_loss.item(), 'policy_adv_abs_max': policy_adv.abs().max().item(), 'approx_kl': ppo_info.approx_kl, 'clipfrac': ppo_info.clipfrac, } def _state_dict_learn(self) -> Dict[str, Any]: """ Overview: Return the state_dict of learn mode, usually including model and optimizer. Returns: - state_dict (:obj:`Dict[str, Any]`): the dict of current policy learn state, for saving and restoring. """ return { 'model': self._learn_model.state_dict(), 'optimizer_ac': self._optimizer_ac.state_dict(), 'optimizer_aux_critic': self._optimizer_aux_critic.state_dict(), } def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: """ Overview: Load the state_dict variable into policy learn mode. Arguments: - state_dict (:obj:`Dict[str, Any]`): the dict of policy learn state saved before.\ When the value is distilled into the policy network, we need to make sure the policy \ network does not change the action predictions, we need two optimizers, \ _optimizer_ac is used in policy net, and _optimizer_aux_critic is used in value net. .. tip:: If you want to only load some parts of model, you can simply set the ``strict`` argument in \ load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \ complicated operation. """ self._learn_model.load_state_dict(state_dict['model']) self._optimizer_ac.load_state_dict(state_dict['optimizer_ac']) self._optimizer_aux_critic.load_state_dict(state_dict['optimizer_aux_critic']) def _init_collect(self) -> None: """ Overview: Initialize the collect mode of policy, including related attributes and modules. For PPG, 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``. """ self._unroll_len = self._cfg.collect.unroll_len self._collect_model = model_wrap(self._model, wrapper_name='multinomial_sample') # TODO continuous action space exploration self._collect_model.reset() self._gamma = self._cfg.collect.discount_factor self._gae_lambda = self._cfg.collect.gae_lambda def _forward_collect(self, data: dict) -> 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[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 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 PPGPolicy: ``ding.policy.tests.test_ppg``. """ 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, mode='compute_actor_critic') 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 PPG, it contains obs, next_obs, action, reward, done, logit, value. Arguments: - obs (:obj:`Any`): Env observation - model_output (:obj:`dict`): The output of the policy network with the observation \ as input. For PPG, 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`): 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, 'logit': model_output['logit'], 'action': model_output['action'], 'value': model_output['value'], 'reward': timestep.reward, 'done': timestep.done, } return transition def _get_train_sample(self, data: List[Dict[str, Any]]) -> 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 PPG, a train sample is a processed transition with new computed \ ``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: - data (: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:`dict`): 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) 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_value = self._collect_model.forward( data[-1]['next_obs'].unsqueeze(0), mode='compute_actor_critic' )['value'] 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 return get_train_sample(data, self._unroll_len) def _get_batch_size(self) -> Dict[str, int]: """ Overview: Get learn batch size. In the PPG algorithm, different networks require different data.\ We need to get data['policy'] and data['value'] to train policy net and value net,\ this function is used to get the batch size of data['policy'] and data['value']. Returns: - output (:obj:`dict[str, int]`): Dict type data, including str type batch size and int type batch size. """ bs = self._cfg.learn.batch_size return {'policy': bs, 'value': bs} def _init_eval(self) -> None: """ Overview: Initialize the eval mode of policy, including related attributes and modules. For PPG, 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``. """ 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 PPG 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[str, 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 PPGPolicy: ``ding.policy.tests.test_ppg``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) if self._cuda: data = to_device(data, self._device) self._eval_model.eval() with torch.no_grad(): output = self._eval_model.forward(data, 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 _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: - vars (:obj:`List[str]`): The list of the necessary keys to be logged. """ return [ 'policy_cur_lr', 'value_cur_lr', 'policy_loss', 'value_loss', 'entropy_loss', 'policy_adv_abs_max', 'approx_kl', 'clipfrac', 'aux_value_loss', 'auxiliary_loss', 'behavioral_cloning_loss', ] def learn_aux(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Overview: The auxiliary phase training, where the value is distilled into the policy network. In PPG algorithm, \ we use the value function loss as the auxiliary objective, thereby sharing features between the policy \ and value function while minimizing distortions to the policy. We also use behavioral cloning loss to \ optimize the auxiliary objective while otherwise preserving the original policy. Returns: - aux_loss (:obj:`Tuple[torch.Tensor, torch.Tensor, torch.Tensor]`): Including average auxiliary loss\ average behavioral cloning loss, and average auxiliary value loss. """ aux_memories = self._aux_memories # gather states and target values into one tensor data = {} states = [] actions = [] return_ = [] old_values = [] weights = [] for memory in aux_memories: # for memory in memories: states.append(memory['obs']) actions.append(memory['action']) return_.append(memory['return_']) old_values.append(memory['value']) if memory['weight'] is None: weight = torch.ones_like(memory['action']) else: weight = torch.tensor(memory['weight']) weights.append(weight) data['obs'] = torch.cat(states) data['action'] = torch.cat(actions) data['return_'] = torch.cat(return_) data['value'] = torch.cat(old_values) data['weight'] = torch.cat(weights).float() # compute current policy logit_old with torch.no_grad(): data['logit_old'] = self._model.forward(data['obs'], mode='compute_actor')['logit'] # prepared dataloader for auxiliary phase training dl = create_shuffled_dataloader(data, self._cfg.learn.batch_size) # the proposed auxiliary phase training # where the value is distilled into the policy network, # while making sure the policy network does not change the action predictions (kl div loss) i = 0 auxiliary_loss_ = 0 behavioral_cloning_loss_ = 0 value_loss_ = 0 for epoch in range(self._aux_train_epoch): for data in dl: policy_output = self._model.forward(data['obs'], mode='compute_actor_critic') # Calculate ppg error 'logit_new', 'logit_old', 'action', 'value_new', 'value_old', 'return_', 'weight' data_ppg = ppg_data( policy_output['logit'], data['logit_old'], data['action'], policy_output['value'], data['value'], data['return_'], data['weight'] ) ppg_joint_loss = ppg_joint_error(data_ppg, self._clip_ratio) wb = self._aux_bc_weight total_loss = ppg_joint_loss.auxiliary_loss + wb * ppg_joint_loss.behavioral_cloning_loss # # policy network loss copmoses of both the kl div loss as well as the auxiliary loss # aux_loss = clipped_value_loss(policy_values, rewards, old_values, self.value_clip) # loss_kl = F.kl_div(action_logprobs, old_action_probs, reduction='batchmean') # policy_loss = aux_loss + loss_kl self._optimizer_ac.zero_grad() total_loss.backward() self._optimizer_ac.step() # paper says it is important to train the value network extra during the auxiliary phase # Calculate ppg error 'value_new', 'value_old', 'return_', 'weight' values = self._model.forward(data['obs'], mode='compute_critic')['value'] data_aux = ppo_value_data(values, data['value'], data['return_'], data['weight']) value_loss = ppo_value_error(data_aux, self._clip_ratio) self._optimizer_aux_critic.zero_grad() value_loss.backward() self._optimizer_aux_critic.step() auxiliary_loss_ += ppg_joint_loss.auxiliary_loss.item() behavioral_cloning_loss_ += ppg_joint_loss.behavioral_cloning_loss.item() value_loss_ += value_loss.item() i += 1 self._aux_memories = [] return auxiliary_loss_ / i, behavioral_cloning_loss_ / i, value_loss_ / i @POLICY_REGISTRY.register('ppg_offpolicy') class PPGOffPolicy(Policy): """ Overview: Policy class of PPG algorithm with off-policy training mode. Off-policy PPG contains two different data \ max_use buffers. The policy buffer offers data for policy phase , while the value buffer provides auxiliary \ phase's data. The whole training procedure is similar to off-policy PPO but execute additional auxiliary \ phase with a fixed frequency. Interface: ``_init_learn``, ``_data_preprocess_learn``, ``_forward_learn``, ``_state_dict_learn``, \ ``_load_state_dict_learn``, ``_init_collect``, ``_forward_collect``, ``_process_transition``, \ ``_get_train_sample``, ``_get_batch_size``, ``_init_eval``, ``_forward_eval``, ``default_model``, \ ``_monitor_vars_learn``, ``learn_aux``. Config: == ==================== ======== ============== ======================================== ======================= ID Symbol Type Default Value Description Other(Shape) == ==================== ======== ============== ======================================== ======================= 1 ``type`` str ppg | 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 5 | 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.value_`` float 1.0 | The loss weight of value network | policy network weight | ``weight`` | is set to 1 8 | ``learn.entropy_`` float 0.01 | The loss weight of entropy | policy network weight | ``weight`` | regularization | is set to 1 9 | ``learn.clip_`` float 0.2 | PPO clip ratio | ``ratio`` 10 | ``learn.adv_`` bool False | Whether to use advantage norm in | ``norm`` | a whole training batch 11 | ``learn.aux_`` int 5 | The frequency(normal update times) | ``freq`` | of auxiliary phase training 12 | ``learn.aux_`` int 6 | The training epochs of auxiliary | ``train_epoch`` | phase 13 | ``learn.aux_`` int 1 | The loss weight of behavioral_cloning | ``bc_weight`` | in auxiliary phase 14 | ``collect.dis`` float 0.99 | Reward's future discount factor, aka. | may be 1 when sparse | ``count_factor`` | gamma | reward env 15 | ``collect.gae_`` float 0.95 | GAE lambda factor for the balance | ``lambda`` | of bias and variance(1-step td and mc) == ==================== ======== ============== ======================================== ======================= """ config = dict( # (str) RL policy register name (refer to function "POLICY_REGISTRY"). type='ppg_offpolicy', # (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=False, priority=False, # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True. priority_IS_weight=False, # (bool) Whether to need policy data in process transition transition_with_policy_data=True, learn=dict( update_per_collect=5, batch_size=64, learning_rate=0.001, # ============================================================== # 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.01, # (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=False, # (int) The frequency(normal update times) of auxiliary phase training aux_freq=5, # (int) The training epochs of auxiliary phase aux_train_epoch=6, # (int) The loss weight of behavioral_cloning in auxiliary phase aux_bc_weight=1, ignore_done=False, ), collect=dict( # n_sample=64, 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(), other=dict( replay_buffer=dict( # PPG use two separate buffer for different reuse multi_buffer=True, policy=dict(replay_buffer_size=1000, ), value=dict(replay_buffer_size=1000, ), ), ), ) 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. """ return 'ppg', ['ding.model.template.ppg'] def _init_learn(self) -> None: """ Overview: Initialize the learn mode of policy, including related attributes and modules. For PPG, it mainly \ contains optimizer, algorithm-specific arguments such as aux_bc_weight and aux_train_epoch. 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``. """ # Optimizer self._optimizer_ac = Adam(self._model.actor_critic.parameters(), lr=self._cfg.learn.learning_rate) self._optimizer_aux_critic = Adam(self._model.aux_critic.parameters(), lr=self._cfg.learn.learning_rate) self._learn_model = model_wrap(self._model, wrapper_name='base') # Algorithm config 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 PPG" 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 # Main model self._learn_model.reset() # Auxiliary memories self._aux_train_epoch = self._cfg.learn.aux_train_epoch self._train_iteration = 0 self._aux_memories = [] self._aux_bc_weight = self._cfg.learn.aux_bc_weight def _data_preprocess_learn(self, data: List[Any]) -> dict: """ Overview: Preprocess the data to fit the required data format for learning, including \ collate(stack data into batch), ignore done(in some fake terminate env),\ prepare loss weight per training sample, and cpu tensor to cuda. Arguments: - data (:obj:`List[Dict[str, Any]]`): The data collected from collect function. Returns: - data (:obj:`Dict[str, Any]`): The processed data, including at least ['done', 'weight']. """ # data preprocess for k, data_item in data.items(): data_item = default_collate(data_item) ignore_done = self._cfg.learn.ignore_done if ignore_done: data_item['done'] = None else: data_item['done'] = data_item['done'].float() data_item['weight'] = None data[k] = data_item if self._cuda: data = to_device(data, self._device) return data def _forward_learn(self, data: dict) -> Dict[str, Any]: """ Overview: Forward and backward function of learn mode. Arguments: - data (:obj:`Dict[str, Any]`): Input data used for policy forward, including the \ collected training samples from replay buffer. For each element in dict, the key of the \ dict is the name of data items and the value is the corresponding data. Usually, \ the class type of value is either torch.Tensor or np.ndarray, or a dict/list containing \ either torch.Tensor or np.ndarray items 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 PPGOff, 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: - 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. \ For the detailed definition of the dict, refer to the code of ``_monitor_vars_learn`` method. ReturnsKeys: - necessary: "current lr", "total_loss", "policy_loss", "value_loss", "entropy_loss", \ "adv_abs_max", "approx_kl", "clipfrac", \ "aux_value_loss", "auxiliary_loss", "behavioral_cloning_loss". - current_lr (:obj:`float`): Current learning rate. - total_loss (:obj:`float`): The calculated loss. - policy_loss (:obj:`float`): The policy(actor) loss of ppg. - value_loss (:obj:`float`): The value(critic) loss of ppg. - entropy_loss (:obj:`float`): The entropy loss. - auxiliary_loss (:obj:`float`): The auxiliary loss, we use the value function loss \ as the auxiliary objective, thereby sharing features between the policy and value function\ while minimizing distortions to the policy. - aux_value_loss (:obj:`float`): The auxiliary value loss, we need to train the value network extra \ during the auxiliary phase, it's the value loss we train the value network during auxiliary phase. - behavioral_cloning_loss (:obj:`float`): The behavioral cloning loss, used to optimize the auxiliary\ objective while otherwise preserving the original policy. """ data = self._data_preprocess_learn(data) # ==================== # PPG forward # ==================== self._learn_model.train() policy_data, value_data = data['policy'], data['value'] policy_adv, value_adv = policy_data['adv'], value_data['adv'] return_ = value_data['value'] + value_adv if self._adv_norm: # Normalize advantage in a total train_batch policy_adv = (policy_adv - policy_adv.mean()) / (policy_adv.std() + 1e-8) value_adv = (value_adv - value_adv.mean()) / (value_adv.std() + 1e-8) # Policy Phase(Policy) policy_output = self._learn_model.forward(policy_data['obs'], mode='compute_actor') policy_error_data = ppo_policy_data( policy_output['logit'], policy_data['logit'], policy_data['action'], policy_adv, policy_data['weight'] ) ppo_policy_loss, ppo_info = ppo_policy_error(policy_error_data, self._clip_ratio) policy_loss = ppo_policy_loss.policy_loss - self._entropy_weight * ppo_policy_loss.entropy_loss self._optimizer_ac.zero_grad() policy_loss.backward() self._optimizer_ac.step() # Policy Phase(Value) value_output = self._learn_model.forward(value_data['obs'], mode='compute_critic') value_error_data = ppo_value_data(value_output['value'], value_data['value'], return_, value_data['weight']) value_loss = self._value_weight * ppo_value_error(value_error_data, self._clip_ratio) self._optimizer_aux_critic.zero_grad() value_loss.backward() self._optimizer_aux_critic.step() # ==================== # PPG update # use aux loss after iterations and reset aux_memories # ==================== # Auxiliary Phase # record data for auxiliary head data = data['value'] data['return_'] = return_.data self._aux_memories.append(copy.deepcopy(data)) self._train_iteration += 1 total_loss = policy_loss + value_loss if self._train_iteration % self._cfg.learn.aux_freq == 0: aux_loss, bc_loss, aux_value_loss = self.learn_aux() total_loss += aux_loss + bc_loss + aux_value_loss return { 'policy_cur_lr': self._optimizer_ac.defaults['lr'], 'value_cur_lr': self._optimizer_aux_critic.defaults['lr'], 'policy_loss': ppo_policy_loss.policy_loss.item(), 'value_loss': value_loss.item(), 'entropy_loss': ppo_policy_loss.entropy_loss.item(), 'policy_adv_abs_max': policy_adv.abs().max().item(), 'approx_kl': ppo_info.approx_kl, 'clipfrac': ppo_info.clipfrac, 'aux_value_loss': aux_value_loss, 'auxiliary_loss': aux_loss, 'behavioral_cloning_loss': bc_loss, 'total_loss': total_loss.item(), } else: return { 'policy_cur_lr': self._optimizer_ac.defaults['lr'], 'value_cur_lr': self._optimizer_aux_critic.defaults['lr'], 'policy_loss': ppo_policy_loss.policy_loss.item(), 'value_loss': value_loss.item(), 'entropy_loss': ppo_policy_loss.entropy_loss.item(), 'policy_adv_abs_max': policy_adv.abs().max().item(), 'approx_kl': ppo_info.approx_kl, 'clipfrac': ppo_info.clipfrac, 'total_loss': total_loss.item(), } def _state_dict_learn(self) -> Dict[str, Any]: """ Overview: Return the state_dict of learn mode, usually including model and optimizer. Returns: - state_dict (:obj:`Dict[str, Any]`): the dict of current policy learn state, for saving and restoring. """ return { 'model': self._learn_model.state_dict(), 'optimizer_ac': self._optimizer_ac.state_dict(), 'optimizer_aux_critic': self._optimizer_aux_critic.state_dict(), } def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None: """ Overview: Load the state_dict variable into policy learn mode. Arguments: - state_dict (:obj:`Dict[str, Any]`): the dict of policy learn state saved before.\ When the value is distilled into the policy network, we need to make sure the policy \ network does not change the action predictions, we need two optimizers, \ _optimizer_ac is used in policy net, and _optimizer_aux_critic is used in value net. .. tip:: If you want to only load some parts of model, you can simply set the ``strict`` argument in \ load_state_dict to ``False``, or refer to ``ding.torch_utils.checkpoint_helper`` for more \ complicated operation. """ self._learn_model.load_state_dict(state_dict['model']) self._optimizer_ac.load_state_dict(state_dict['optimizer_ac']) self._optimizer_aux_critic.load_state_dict(state_dict['optimizer_aux_critic']) def _init_collect(self) -> None: """ Overview: Initialize the collect mode of policy, including related attributes and modules. For PPO, 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``. """ self._unroll_len = self._cfg.collect.unroll_len self._collect_model = model_wrap(self._model, wrapper_name='multinomial_sample') # TODO continuous action space exploration self._collect_model.reset() self._gamma = self._cfg.collect.discount_factor self._gae_lambda = self._cfg.collect.gae_lambda def _forward_collect(self, data: dict) -> 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[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 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 PPGOffPolicy: ``ding.policy.tests.test_ppg``. """ 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, mode='compute_actor_critic') 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 PPG, it contains obs, next_obs, action, reward, done, logit, value. Arguments: - obs (:obj:`Any`): Env observation - model_output (:obj:`dict`): The output of the policy network with the observation \ as input. For PPG, 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`): 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, 'logit': model_output['logit'], 'action': model_output['action'], '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 PPG, a train sample is a processed transition with new computed \ ``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: - data (:obj:`list`): 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:`dict`): 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 = get_gae_with_default_last_value( data, data[-1]['done'], gamma=self._gamma, gae_lambda=self._gae_lambda, cuda=False, ) data = get_train_sample(data, self._unroll_len) for d in data: d['buffer_name'] = ["policy", "value"] return data def _get_batch_size(self) -> Dict[str, int]: """ Overview: Get learn batch size. In the PPG algorithm, different networks require different data.\ We need to get data['policy'] and data['value'] to train policy net and value net,\ this function is used to get the batch size of data['policy'] and data['value']. Returns: - output (:obj:`dict[str, int]`): Dict type data, including str type batch size and int type batch size. """ bs = self._cfg.learn.batch_size return {'policy': bs, 'value': bs} def _init_eval(self) -> None: """ Overview: Initialize the eval mode of policy, including related attributes and modules. For PPG, 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``. """ self._eval_model = model_wrap(self._model, wrapper_name='argmax_sample') self._eval_model.reset() def _forward_eval(self, data: dict) -> dict: r""" 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 PPG 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[str, 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 PPGOffPolicy: ``ding.policy.tests.test_ppg``. """ data_id = list(data.keys()) data = default_collate(list(data.values())) if self._cuda: data = to_device(data, self._device) self._eval_model.eval() with torch.no_grad(): output = self._eval_model.forward(data, 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 _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: - vars (:obj:`List[str]`): The list of the necessary keys to be logged. """ return [ 'policy_cur_lr', 'value_cur_lr', 'policy_loss', 'value_loss', 'entropy_loss', 'policy_adv_abs_max', 'approx_kl', 'clipfrac', 'aux_value_loss', 'auxiliary_loss', 'behavioral_cloning_loss', ] def learn_aux(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Overview: The auxiliary phase training, where the value is distilled into the policy network. In PPG algorithm, \ we use the value function loss as the auxiliary objective, thereby sharing features between the policy \ and value function while minimizing distortions to the policy. We also use behavioral cloning loss to \ optimize the auxiliary objective while otherwise preserving the original policy. Returns: - aux_loss (:obj:`Tuple[torch.Tensor, torch.Tensor, torch.Tensor]`): Including average auxiliary loss\ average behavioral cloning loss, and average auxiliary value loss. """ aux_memories = self._aux_memories # gather states and target values into one tensor data = {} states = [] actions = [] return_ = [] old_values = [] weights = [] for memory in aux_memories: # for memory in memories: states.append(memory['obs']) actions.append(memory['action']) return_.append(memory['return_']) old_values.append(memory['value']) if memory['weight'] is None: weight = torch.ones_like(memory['action']) else: weight = torch.tensor(memory['weight']) weights.append(weight) data['obs'] = torch.cat(states) data['action'] = torch.cat(actions) data['return_'] = torch.cat(return_) data['value'] = torch.cat(old_values) data['weight'] = torch.cat(weights) # compute current policy logit_old with torch.no_grad(): data['logit_old'] = self._model.forward(data['obs'], mode='compute_actor')['logit'] # prepared dataloader for auxiliary phase training dl = create_shuffled_dataloader(data, self._cfg.learn.batch_size) # the proposed auxiliary phase training # where the value is distilled into the policy network, # while making sure the policy network does not change the action predictions (kl div loss) i = 0 auxiliary_loss_ = 0 behavioral_cloning_loss_ = 0 value_loss_ = 0 for epoch in range(self._aux_train_epoch): for data in dl: policy_output = self._model.forward(data['obs'], mode='compute_actor_critic') # Calculate ppg error 'logit_new', 'logit_old', 'action', 'value_new', 'value_old', 'return_', 'weight' data_ppg = ppg_data( policy_output['logit'], data['logit_old'], data['action'], policy_output['value'], data['value'], data['return_'], data['weight'] ) ppg_joint_loss = ppg_joint_error(data_ppg, self._clip_ratio) wb = self._aux_bc_weight total_loss = ppg_joint_loss.auxiliary_loss + wb * ppg_joint_loss.behavioral_cloning_loss # # policy network loss copmoses of both the kl div loss as well as the auxiliary loss # aux_loss = clipped_value_loss(policy_values, rewards, old_values, self.value_clip) # loss_kl = F.kl_div(action_logprobs, old_action_probs, reduction='batchmean') # policy_loss = aux_loss + loss_kl self._optimizer_ac.zero_grad() total_loss.backward() self._optimizer_ac.step() # paper says it is important to train the value network extra during the auxiliary phase # Calculate ppg error 'value_new', 'value_old', 'return_', 'weight' values = self._model.forward(data['obs'], mode='compute_critic')['value'] data_aux = ppo_value_data(values, data['value'], data['return_'], data['weight']) value_loss = ppo_value_error(data_aux, self._clip_ratio) self._optimizer_aux_critic.zero_grad() value_loss.backward() self._optimizer_aux_critic.step() auxiliary_loss_ += ppg_joint_loss.auxiliary_loss.item() behavioral_cloning_loss_ += ppg_joint_loss.behavioral_cloning_loss.item() value_loss_ += value_loss.item() i += 1 self._aux_memories = [] return auxiliary_loss_ / i, behavioral_cloning_loss_ / i, value_loss_ / i