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
[docs]@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