def __init__(self, args, policy, device=torch.device("cpu")):
self.args = args
self.device = device
self.tpdv = dict(dtype=torch.float32, device=device)
self.policy = policy
self.clip_param = args.clip_param
self.ppo_epoch = args.ppo_epoch
self.num_mini_batch = args.num_mini_batch
self.data_chunk_length = args.data_chunk_length
self.value_loss_coef = args.value_loss_coef
self.entropy_coef = args.entropy_coef
self.max_grad_norm = args.max_grad_norm
self.huber_delta = args.huber_delta
self._use_recurrent_policy = args.use_recurrent_policy
self._use_naive_recurrent = args.use_naive_recurrent_policy
self._use_max_grad_norm = args.use_max_grad_norm
self._use_clipped_value_loss = args.use_clipped_value_loss
self._use_huber_loss = args.use_huber_loss
self._use_popart = args.use_popart
self._use_value_active_masks = args.use_value_active_masks
self._use_policy_active_masks = args.use_policy_active_masks
self._use_ob = args.use_ob
if self._use_popart:
self.value_normalizer = PopArt(1, device=self.device)
else:
self.value_normalizer = None
class R_MAPPO():
"""
Trainer class for MAPPO to update policies.
:param args: (argparse.Namespace) arguments containing relevant model, policy, and env information.
:param policy: (R_MAPPO_Policy) policy to update.
:param device: (torch.device) specifies the device to run on (cpu/gpu).
"""
def __init__(self, args, policy, device=torch.device("cpu")):
self.args = args
self.device = device
self.tpdv = dict(dtype=torch.float32, device=device)
self.policy = policy
self.clip_param = args.clip_param
self.ppo_epoch = args.ppo_epoch
self.num_mini_batch = args.num_mini_batch
self.data_chunk_length = args.data_chunk_length
self.value_loss_coef = args.value_loss_coef
self.entropy_coef = args.entropy_coef
self.max_grad_norm = args.max_grad_norm
self.huber_delta = args.huber_delta
self._use_recurrent_policy = args.use_recurrent_policy
self._use_naive_recurrent = args.use_naive_recurrent_policy
self._use_max_grad_norm = args.use_max_grad_norm
self._use_clipped_value_loss = args.use_clipped_value_loss
self._use_huber_loss = args.use_huber_loss
self._use_popart = args.use_popart
self._use_value_active_masks = args.use_value_active_masks
self._use_policy_active_masks = args.use_policy_active_masks
self._use_ob = args.use_ob
if self._use_popart:
self.value_normalizer = PopArt(1, device=self.device)
else:
self.value_normalizer = None
def cal_value_loss(self, values, value_preds_batch, return_batch,
active_masks_batch):
"""
Calculate value function loss.
:param values: (torch.Tensor) value function predictions.
:param value_preds_batch: (torch.Tensor) "old" value predictions from data batch (used for value clip loss)
:param return_batch: (torch.Tensor) reward to go returns.
:param active_masks_batch: (torch.Tensor) denotes if agent is active or dead at a given timesep.
:return value_loss: (torch.Tensor) value function loss.
"""
if self._use_popart:
value_pred_clipped = value_preds_batch + (
values - value_preds_batch).clamp(-self.clip_param,
self.clip_param)
error_clipped = self.value_normalizer(
return_batch) - value_pred_clipped
error_original = self.value_normalizer(return_batch) - values
else:
value_pred_clipped = value_preds_batch + (
values - value_preds_batch).clamp(-self.clip_param,
self.clip_param)
error_clipped = return_batch - value_pred_clipped
error_original = return_batch - values
if self._use_huber_loss:
value_loss_clipped = huber_loss(error_clipped, self.huber_delta)
value_loss_original = huber_loss(error_original, self.huber_delta)
else:
value_loss_clipped = mse_loss(error_clipped)
value_loss_original = mse_loss(error_original)
if self._use_clipped_value_loss:
value_loss = torch.max(value_loss_original, value_loss_clipped)
else:
value_loss = value_loss_original
if self._use_value_active_masks:
value_loss = (value_loss *
active_masks_batch).sum() / active_masks_batch.sum()
else:
value_loss = value_loss.mean()
return value_loss
def ppo_update(self, sample, update_actor=True):
"""
Update actor and critic networks.
:param sample: (Tuple) contains data batch with which to update networks.
:update_actor: (bool) whether to update actor network.
:return value_loss: (torch.Tensor) value function loss.
:return critic_grad_norm: (torch.Tensor) gradient norm from critic up9date.
;return policy_loss: (torch.Tensor) actor(policy) loss value.
:return dist_entropy: (torch.Tensor) action entropies.
:return actor_grad_norm: (torch.Tensor) gradient norm from actor update.
:return imp_weights: (torch.Tensor) importance sampling weights.
"""
share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch, \
value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch, \
adv_targ, available_actions_batch = sample
old_action_log_probs_batch = check(old_action_log_probs_batch).to(
**self.tpdv)
adv_targ = check(adv_targ).to(**self.tpdv)
value_preds_batch = check(value_preds_batch).to(**self.tpdv)
return_batch = check(return_batch).to(**self.tpdv)
active_masks_batch = check(active_masks_batch).to(**self.tpdv)
# Reshape to do in a single forward pass for all steps
values, action_log_probs, dist_entropy = self.policy.evaluate_actions(
share_obs_batch, obs_batch, rnn_states_batch,
rnn_states_critic_batch, actions_batch, masks_batch,
available_actions_batch, active_masks_batch)
# actor update
imp_weights = torch.exp(action_log_probs - old_action_log_probs_batch)
surr1 = imp_weights * adv_targ
surr2 = torch.clamp(imp_weights, 1.0 - self.clip_param,
1.0 + self.clip_param) * adv_targ
if self._use_policy_active_masks:
policy_action_loss = (
-torch.sum(torch.min(surr1, surr2), dim=-1, keepdim=True) *
active_masks_batch).sum() / active_masks_batch.sum()
else:
policy_action_loss = -torch.sum(
torch.min(surr1, surr2), dim=-1, keepdim=True).mean()
policy_loss = policy_action_loss
self.policy.actor_optimizer.zero_grad()
if update_actor:
(policy_loss - dist_entropy * self.entropy_coef).backward()
if self._use_max_grad_norm:
actor_grad_norm = nn.utils.clip_grad_norm_(
self.policy.actor.parameters(), self.max_grad_norm)
else:
actor_grad_norm = get_gard_norm(self.policy.actor.parameters())
self.policy.actor_optimizer.step()
# critic update
value_loss = self.cal_value_loss(values, value_preds_batch,
return_batch, active_masks_batch)
if self.args.use_q_head:
self.policy.q_optimizer.zero_grad()
(value_loss * self.value_loss_coef).backward()
if self._use_max_grad_norm:
critic_grad_norm = nn.utils.clip_grad_norm_(
self.policy.q_head.parameters(), self.max_grad_norm)
else:
critic_grad_norm = get_gard_norm(
self.policy.q_head.parameters())
self.policy.q_optimizer.step()
else:
self.policy.critic_optimizer.zero_grad()
(value_loss * self.value_loss_coef).backward()
if self._use_max_grad_norm:
critic_grad_norm = nn.utils.clip_grad_norm_(
self.policy.critic.parameters(), self.max_grad_norm)
else:
critic_grad_norm = get_gard_norm(
self.policy.critic.parameters())
self.policy.critic_optimizer.step()
return value_loss, critic_grad_norm, policy_loss, dist_entropy, actor_grad_norm, imp_weights
def piob_corrected(self, q, pi):
"""
Calculating the optimal baseline
:param q: Q values of all actions under the same states
:param pi: Probability distributions of actions
:return: The optimal baseline
"""
M = torch.norm(pi)**2 + 1
xweight = M - 2 * pi
enum = (pi * xweight * q).sum(-1)
deno = (pi * xweight).sum(-1)
return (enum / deno).unsqueeze(-1).detach().numpy()
def train(self, buffer, update_actor=True):
"""
Perform a training update using minibatch GD.
:param buffer: (SharedReplayBuffer) buffer containing training data.
:param update_actor: (bool) whether to update actor network.
:return train_info: (dict) contains information regarding training update (e.g. loss, grad norms, etc).
"""
if self.args.use_q_head:
if self._use_popart:
# q_vals = torch.Tensor(buffer.q_value_preds[:-1])
# pi = torch.softmax(torch.Tensor(buffer.pi[:-1]), dim=-1).clamp(0.0001, 0.9999)
# ob = self.piob_corrected(q_vals, pi)
# actions = torch.tensor(buffer.actions, dtype=torch.int64)
# q_taken = torch.gather(q_vals, dim=-1, index=actions).numpy()
# advantages = q_taken - buffer.ob[:-1]
if self._use_ob:
advantages = buffer.rewards + self.args.gamma * self.value_normalizer.denormalize(
buffer.value_preds[1:]
) * buffer.masks[1:] - buffer.ob[:-1]
else:
advantages = buffer.returns[:
-1] - self.value_normalizer.denormalize(
buffer.value_preds[:-1])
else:
advantages = buffer.returns[:-1] - buffer.value_preds[:-1]
advantages_copy = advantages.copy()
advantages_copy[buffer.active_masks[:-1] == 0.0] = np.nan
mean_advantages = np.nanmean(advantages_copy)
std_advantages = np.nanstd(advantages_copy)
advantages = (advantages - mean_advantages) / (std_advantages + 1e-5)
train_info = {}
train_info['value_loss'] = 0
train_info['policy_loss'] = 0
train_info['dist_entropy'] = 0
train_info['actor_grad_norm'] = 0
train_info['critic_grad_norm'] = 0
train_info['ratio'] = 0
for _ in range(self.ppo_epoch):
if self._use_recurrent_policy:
data_generator = buffer.recurrent_generator(
advantages, self.num_mini_batch, self.data_chunk_length)
elif self._use_naive_recurrent:
data_generator = buffer.naive_recurrent_generator(
advantages, self.num_mini_batch)
else:
data_generator = buffer.feed_forward_generator(
advantages, self.num_mini_batch)
for sample in data_generator:
value_loss, critic_grad_norm, policy_loss, dist_entropy, actor_grad_norm, imp_weights \
= self.ppo_update(sample, update_actor)
train_info['value_loss'] += value_loss.item()
train_info['policy_loss'] += policy_loss.item()
train_info['dist_entropy'] += dist_entropy.item()
train_info['actor_grad_norm'] += actor_grad_norm
train_info['critic_grad_norm'] += critic_grad_norm
train_info['ratio'] += imp_weights.mean()
num_updates = self.ppo_epoch * self.num_mini_batch
for k in train_info.keys():
train_info[k] /= num_updates
return train_info
def prep_training(self):
self.policy.actor.train()
if self.args.use_q_head:
self.policy.q_head.train()
else:
self.policy.critic.train()
def prep_rollout(self):
self.policy.actor.eval()
if self.args.use_q_head:
self.policy.q_head.eval()
else:
self.policy.critic.eval()
class R_MAPPO():
def __init__(self,
args,
policy,
device=torch.device("cpu")):
self.device = device
self.tpdv = dict(dtype=torch.float32, device=device)
self.policy = policy
self.clip_param = args.clip_param
self.ppo_epoch = args.ppo_epoch
self.num_mini_batch = args.num_mini_batch
self.data_chunk_length = args.data_chunk_length
self.policy_value_loss_coef = args.policy_value_loss_coef
self.value_loss_coef = args.value_loss_coef
self.entropy_coef = args.entropy_coef
self.max_grad_norm = args.max_grad_norm
self.huber_delta = args.huber_delta
self._use_recurrent_policy = args.use_recurrent_policy
self._use_naive_recurrent = args.use_naive_recurrent_policy
self._use_max_grad_norm = args.use_max_grad_norm
self._use_clipped_value_loss = args.use_clipped_value_loss
self._use_huber_loss = args.use_huber_loss
self._use_popart = args.use_popart
self._use_value_active_masks = args.use_value_active_masks
self._use_policy_active_masks = args.use_policy_active_masks
self._use_policy_vhead = args.use_policy_vhead
if self._use_popart:
self.value_normalizer = PopArt(1, device=self.device)
else:
self.value_normalizer = None
def cal_value_loss(self, values, value_preds_batch, return_batch, active_masks_batch):
if self._use_popart:
value_pred_clipped = value_preds_batch + (values - value_preds_batch).clamp(-self.clip_param, self.clip_param)
error_clipped = self.value_normalizer(return_batch) - value_pred_clipped
error_original = self.value_normalizer(return_batch) - values
else:
value_pred_clipped = value_preds_batch + (values - value_preds_batch).clamp(-self.clip_param, self.clip_param)
error_clipped = return_batch - value_pred_clipped
error_original = return_batch - values
if self._use_huber_loss:
value_loss_clipped = huber_loss(error_clipped, self.huber_delta)
value_loss_original = huber_loss(error_original, self.huber_delta)
else:
value_loss_clipped = mse_loss(error_clipped)
value_loss_original = mse_loss(error_original)
if self._use_clipped_value_loss:
value_loss = torch.max(value_loss_original, value_loss_clipped)
else:
value_loss = value_loss_original
if self._use_value_active_masks:
value_loss = (value_loss * active_masks_batch).sum() / active_masks_batch.sum()
else:
value_loss = value_loss.mean()
return value_loss
def ppo_update(self, sample, turn_on=True):
share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch, \
value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch, \
adv_targ, available_actions_batch = sample
old_action_log_probs_batch = check(old_action_log_probs_batch).to(**self.tpdv)
adv_targ = check(adv_targ).to(**self.tpdv)
value_preds_batch = check(value_preds_batch).to(**self.tpdv)
return_batch = check(return_batch).to(**self.tpdv)
active_masks_batch = check(active_masks_batch).to(**self.tpdv)
# Reshape to do in a single forward pass for all steps
values, action_log_probs, dist_entropy, policy_values = self.policy.evaluate_actions(share_obs_batch,
obs_batch,
rnn_states_batch,
rnn_states_critic_batch,
actions_batch,
masks_batch,
available_actions_batch,
active_masks_batch)
# actor update
ratio = torch.exp(action_log_probs - old_action_log_probs_batch)
surr1 = ratio * adv_targ
surr2 = torch.clamp(ratio, 1.0 - self.clip_param, 1.0 + self.clip_param) * adv_targ
if self._use_policy_active_masks:
policy_action_loss = (-torch.sum(torch.min(surr1, surr2), dim=-1, keepdim=True) * active_masks_batch).sum() / active_masks_batch.sum()
else:
policy_action_loss = -torch.sum(torch.min(surr1, surr2), dim=-1, keepdim=True).mean()
if self._use_policy_vhead:
policy_value_loss = self.cal_value_loss(policy_values, value_preds_batch, return_batch, active_masks_batch)
policy_loss = policy_action_loss + policy_value_loss * self.policy_value_loss_coef
else:
policy_loss = policy_action_loss
self.policy.actor_optimizer.zero_grad()
if turn_on:
(policy_loss - dist_entropy * self.entropy_coef).backward()
if self._use_max_grad_norm:
actor_grad_norm = nn.utils.clip_grad_norm_(self.policy.actor.parameters(), self.max_grad_norm)
else:
actor_grad_norm = get_gard_norm(self.policy.actor.parameters())
self.policy.actor_optimizer.step()
# critic update
value_loss = self.cal_value_loss(values, value_preds_batch, return_batch, active_masks_batch)
self.policy.critic_optimizer.zero_grad()
(value_loss * self.value_loss_coef).backward()
if self._use_max_grad_norm:
critic_grad_norm = nn.utils.clip_grad_norm_(self.policy.critic.parameters(), self.max_grad_norm)
else:
critic_grad_norm = get_gard_norm(self.policy.critic.parameters())
self.policy.critic_optimizer.step()
return value_loss, critic_grad_norm, policy_loss, dist_entropy, actor_grad_norm, ratio
def train(self, buffer, turn_on=True):
if self._use_popart:
advantages = buffer.returns[:-1] - self.value_normalizer.denormalize(buffer.value_preds[:-1])
else:
advantages = buffer.returns[:-1] - buffer.value_preds[:-1]
advantages_copy = advantages.copy()
advantages_copy[buffer.active_masks[:-1] == 0.0] = np.nan
mean_advantages = np.nanmean(advantages_copy)
std_advantages = np.nanstd(advantages_copy)
advantages = (advantages - mean_advantages) / (std_advantages + 1e-5)
train_info = {}
train_info['value_loss'] = 0
train_info['policy_loss'] = 0
train_info['dist_entropy'] = 0
train_info['actor_grad_norm'] = 0
train_info['critic_grad_norm'] = 0
train_info['ratio'] = 0
for _ in range(self.ppo_epoch):
if self._use_recurrent_policy:
data_generator = buffer.recurrent_generator(advantages, self.num_mini_batch, self.data_chunk_length)
elif self._use_naive_recurrent:
data_generator = buffer.naive_recurrent_generator(advantages, self.num_mini_batch)
else:
data_generator = buffer.feed_forward_generator(advantages, self.num_mini_batch)
for sample in data_generator:
value_loss, critic_grad_norm, policy_loss, dist_entropy, actor_grad_norm, ratio \
= self.ppo_update(sample, turn_on)
train_info['value_loss'] += value_loss.item()
train_info['policy_loss'] += policy_loss.item()
train_info['dist_entropy'] += dist_entropy.item()
train_info['actor_grad_norm'] += actor_grad_norm
train_info['critic_grad_norm'] += critic_grad_norm
train_info['ratio'] += ratio.mean()
num_updates = self.ppo_epoch * self.num_mini_batch
for k in train_info.keys():
train_info[k] /= num_updates
return train_info
def prep_training(self):
self.policy.actor.train()
self.policy.critic.train()
def prep_rollout(self):
self.policy.actor.eval()
self.policy.critic.eval()
class R_MAPPG():
def __init__(self, args, policy, device=torch.device("cpu")):
self.device = device
self.tpdv = dict(dtype=torch.float32, device=device)
self.policy = policy
self.clip_param = args.clip_param
self.ppo_epoch = args.ppo_epoch
self.aux_epoch = args.aux_epoch
self.num_mini_batch = args.num_mini_batch
self.data_chunk_length = args.data_chunk_length
self.value_loss_coef = args.value_loss_coef
self.entropy_coef = args.entropy_coef
self.clone_coef = args.clone_coef
self.max_grad_norm = args.max_grad_norm
self.huber_delta = args.huber_delta
self._use_recurrent_policy = args.use_recurrent_policy
self._use_naive_recurrent = args.use_naive_recurrent_policy
self._use_max_grad_norm = args.use_max_grad_norm
self._use_clipped_value_loss = args.use_clipped_value_loss
self._use_huber_loss = args.use_huber_loss
self._use_popart = args.use_popart
self._use_value_active_masks = args.use_value_active_masks
if self._use_popart:
self.value_normalizer = PopArt(1, device=self.device)
else:
self.value_normalizer = None
def policy_loss_update(self, sample):
share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch, \
value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch, \
adv_targ, available_actions_batch = sample
old_action_log_probs_batch = check(old_action_log_probs_batch).to(
**self.tpdv)
adv_targ = check(adv_targ).to(**self.tpdv)
active_masks_batch = check(active_masks_batch).to(**self.tpdv)
# Reshape to do in a single forward pass for all steps
action_log_probs, dist_entropy = self.policy.evaluate_actions(
obs_batch, rnn_states_batch, actions_batch, masks_batch,
available_actions_batch, active_masks_batch)
ratio = torch.exp(action_log_probs - old_action_log_probs_batch)
surr1 = ratio * adv_targ
surr2 = torch.clamp(ratio, 1.0 - self.clip_param,
1.0 + self.clip_param) * adv_targ
action_loss = (
-torch.sum(torch.min(surr1, surr2), dim=-1, keepdim=True) *
active_masks_batch).sum() / active_masks_batch.sum()
# update common and action network
self.policy.actor_optimizer.zero_grad()
(action_loss - dist_entropy * self.entropy_coef).backward()
if self._use_max_grad_norm:
grad_norm = nn.utils.clip_grad_norm_(
self.policy.actor.parameters(), self.max_grad_norm)
else:
grad_norm = get_gard_norm(self.policy.actor.parameters())
self.policy.actor_optimizer.step()
return action_loss, dist_entropy, grad_norm, ratio
def value_loss_update(self, sample):
share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch, \
value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch, \
adv_targ, available_actions_batch = sample
value_preds_batch = check(value_preds_batch).to(**self.tpdv)
return_batch = check(return_batch).to(**self.tpdv)
active_masks_batch = check(active_masks_batch).to(**self.tpdv)
values = self.policy.get_values(share_obs_batch,
rnn_states_critic_batch, masks_batch)
if self._use_popart:
value_pred_clipped = value_preds_batch + (
values - value_preds_batch).clamp(-self.clip_param,
self.clip_param)
error_clipped = self.value_normalizer(
return_batch) - value_pred_clipped
error_original = self.value_normalizer(return_batch) - values
else:
value_pred_clipped = value_preds_batch + (
values - value_preds_batch).clamp(-self.clip_param,
self.clip_param)
error_clipped = return_batch - value_pred_clipped
error_original = return_batch - values
if self._use_huber_loss:
value_loss_clipped = huber_loss(error_clipped, self.huber_delta)
value_loss_original = huber_loss(error_original, self.huber_delta)
else:
value_loss_clipped = mse_loss(error_clipped)
value_loss_original = mse_loss(error_original)
if self._use_clipped_value_loss:
value_loss = torch.max(value_loss_original, value_loss_clipped)
else:
value_loss = value_loss_original
if self._use_value_active_masks:
value_loss = (value_loss *
active_masks_batch).sum() / active_masks_batch.sum()
else:
value_loss = value_loss.mean()
self.policy.critic_optimizer.zero_grad()
(value_loss * self.value_loss_coef).backward()
if self._use_max_grad_norm:
grad_norm = nn.utils.clip_grad_norm_(
self.policy.critic.parameters(), self.max_grad_norm)
else:
grad_norm = get_gard_norm(self.policy.critic.parameters())
self.policy.critic_optimizer.step()
return value_loss, grad_norm
def update_action_probs(self, buffer):
action_probs = []
for step in range(buffer.episode_length):
if buffer.available_actions is not None:
avail_actions = np.concatenate(buffer.available_actions[step])
else:
avail_actions = None
action_prob = self.policy.get_probs(
np.concatenate(buffer.obs[step]),
np.concatenate(buffer.rnn_states[step]),
np.concatenate(buffer.masks[step]), avail_actions)
action_prob = np.array(
np.split(action_prob.detach().cpu().numpy(),
buffer.n_rollout_threads))
action_probs.append(action_prob)
return np.array(action_probs)
def auxiliary_loss_update(self, sample):
share_obs_batch, obs_batch, rnn_states_batch, rnn_states_critic_batch, actions_batch, \
value_preds_batch, return_batch, masks_batch, active_masks_batch, old_action_log_probs_batch, \
old_action_probs_batch, available_actions_batch = sample
old_action_probs_batch = check(old_action_probs_batch).to(**self.tpdv)
active_masks_batch = check(active_masks_batch).to(**self.tpdv)
value_preds_batch = check(value_preds_batch).to(**self.tpdv)
return_batch = check(return_batch).to(**self.tpdv)
# Reshape to do in a single forward pass for all steps
values, new_action_probs = self.policy.get_policy_values_and_probs(
obs_batch, rnn_states_batch, masks_batch, available_actions_batch)
# kl = sum p * log(p / q) = sum p*(logp-logq) = sum plogp - plogq
# cross-entropy = sum -plogq
eps = (old_action_probs_batch == 0) * 1e-8
old_action_log_probs_batch = torch.log(old_action_probs_batch +
eps.float().detach())
eps = (new_action_probs == 0) * 1e-8
new_action_log_probs = torch.log(new_action_probs +
eps.float().detach())
kl_divergence = torch.sum(
(old_action_probs_batch *
(old_action_log_probs_batch - new_action_log_probs)),
dim=-1,
keepdim=True)
kl_loss = (kl_divergence *
active_masks_batch).sum() / active_masks_batch.sum()
if self._use_popart:
value_pred_clipped = value_preds_batch + (
values - value_preds_batch).clamp(-self.clip_param,
self.clip_param)
error_clipped = self.value_normalizer(
return_batch) - value_pred_clipped
error_original = self.value_normalizer(return_batch) - values
else:
value_pred_clipped = value_preds_batch + (
values - value_preds_batch).clamp(-self.clip_param,
self.clip_param)
error_clipped = return_batch - value_pred_clipped
error_original = return_batch - values
if self._use_huber_loss:
value_loss_clipped = huber_loss(error_clipped, self.huber_delta)
value_loss_original = huber_loss(error_original, self.huber_delta)
else:
value_loss_clipped = mse_loss(error_clipped)
value_loss_original = mse_loss(error_original)
if self._use_clipped_value_loss:
value_loss = torch.max(value_loss_original, value_loss_clipped)
else:
value_loss = value_loss_original
if self._use_value_active_masks:
value_loss = (value_loss *
active_masks_batch).sum() / active_masks_batch.sum()
else:
value_loss = value_loss.mean()
joint_loss = value_loss + self.clone_coef * kl_loss
self.policy.actor_optimizer.zero_grad()
joint_loss.backward()
if self._use_max_grad_norm:
grad_norm = nn.utils.clip_grad_norm_(
self.policy.actor.parameters(), self.max_grad_norm)
else:
grad_norm = get_gard_norm(self.policy.actor.parameters())
self.policy.actor_optimizer.step()
return joint_loss, grad_norm
def train(self, buffer):
if self._use_popart:
advantages = buffer.returns[:
-1] - self.value_normalizer.denormalize(
buffer.value_preds[:-1])
else:
advantages = buffer.returns[:-1] - buffer.value_preds[:-1]
advantages_copy = advantages.copy()
advantages_copy[buffer.active_masks[:-1] == 0.0] = np.nan
mean_advantages = np.nanmean(advantages_copy)
std_advantages = np.nanstd(advantages_copy)
advantages = (advantages - mean_advantages) / (std_advantages + 1e-5)
train_info = {}
train_info['value_loss'] = 0
train_info['action_loss'] = 0
train_info['joint_loss'] = 0
train_info['dist_entropy'] = 0
train_info['actor_grad_norm'] = 0
train_info['critic_grad_norm'] = 0
train_info['joint_grad_norm'] = 0
train_info['ratio'] = 0
# policy phase
for _ in range(self.ppo_epoch):
if self._use_recurrent_policy:
data_generator = buffer.recurrent_generator(
advantages, self.num_mini_batch, self.data_chunk_length)
elif self._use_naive_recurrent_policy:
data_generator = buffer.naive_recurrent_generator(
advantages, self.num_mini_batch)
else:
data_generator = buffer.feed_forward_generator(
advantages, self.num_mini_batch)
for sample in data_generator:
action_loss, dist_entropy, actor_grad_norm, ratio = self.policy_loss_update(
sample)
train_info['action_loss'] += action_loss.item()
train_info['dist_entropy'] += dist_entropy.item()
train_info['actor_grad_norm'] += actor_grad_norm
train_info['ratio'] += ratio.mean()
value_loss, critic_grad_norm = self.value_loss_update(sample)
train_info['value_loss'] += value_loss.item()
train_info['critic_grad_norm'] += critic_grad_norm
# auxiliary phase
action_probs = self.update_action_probs(buffer)
for _ in range(self.aux_epoch):
if self._use_recurrent_policy:
data_generator = buffer.recurrent_generator(
action_probs, self.num_mini_batch, self.data_chunk_length)
elif self._use_naive_recurrent_policy:
data_generator = buffer.naive_recurrent_generator(
action_probs, self.num_mini_batch)
else:
data_generator = buffer.feed_forward_generator(
action_probs, self.num_mini_batch)
# 2. update auxiliary
for sample in data_generator:
joint_loss, joint_grad_norm = self.auxiliary_loss_update(
sample)
train_info['joint_loss'] += joint_loss.item()
train_info['joint_grad_norm'] += joint_grad_norm
value_loss, critic_grad_norm = self.value_loss_update(sample)
train_info['value_loss'] += value_loss.item()
train_info['critic_grad_norm'] += critic_grad_norm
for k in train_info.keys():
if k in [
"action_loss", "actor_grad_norm", "ratio", "dist_entropy"
]:
num_updates = self.ppo_epoch * self.num_mini_batch
train_info[k] /= num_updates
elif k in ["value_loss", "critic_grad_norm"]:
num_updates = (self.ppo_epoch +
self.aux_epoch) * self.num_mini_batch
train_info[k] /= num_updates
elif k in ["joint_loss", "joint_grad_norm"]:
num_updates = self.aux_epoch * self.num_mini_batch
train_info[k] /= num_updates
else:
raise NotImplementedError
return train_info
def prep_training(self):
self.policy.actor.train()
self.policy.critic.train()
def prep_rollout(self):
self.policy.actor.eval()
self.policy.critic.eval()