def _categorical_loss(self, states_t, actions_t, rewards_t, states_tp1,
done_t):
gammas = self._gammas**self._n_step
# actor loss
# For now we have the same actor for all heads of the critic
logits_tp0 = self.critic(states_t, self.actor(states_t))
probs_tp0 = torch.softmax(logits_tp0, dim=-1)
q_values_tp0 = torch.sum(probs_tp0 * self.z, dim=-1)
policy_loss = -torch.mean(q_values_tp0)
# critic loss (kl-divergence between categorical distributions)
done_t = done_t[:, None, :]
# B x 1 x 1
rewards_t = rewards_t[:, None, :]
# B x 1 x 1
gammas = gammas[None, :, None]
# 1 x num_heads x 1
logits_t = self.critic(states_t, actions_t)
# B x num_heads x num_atoms
logits_tp1 = self.target_critic(
states_tp1, self.target_actor(states_tp1)).detach()
# B x num_heads x num_atoms
atoms_target_t = rewards_t + (1 - done_t) * gammas * self.z
# B x num_heads x num_atoms
value_loss = categorical_loss(logits_t.view(-1, self.num_atoms),
logits_tp1.view(-1, self.num_atoms),
atoms_target_t.view(-1, self.num_atoms),
self.z, self.delta_z, self.v_min,
self.v_max)
return policy_loss, value_loss
def _categorical_loss(
self, states_t, actions_t, rewards_t, states_tp1, done_t
):
# actor loss
actions_tp0, log_pi_tp0 = self.actor(states_t, with_log_pi=True)
log_pi_tp0 = log_pi_tp0 / self.reward_scale
logits_tp0 = [x(states_t, actions_tp0) for x in self.critics]
probs_tp0 = [F.softmax(x, dim=-1) for x in logits_tp0]
q_values_tp0 = [
torch.sum(x * self.z, dim=-1).unsqueeze_(-1) for x in probs_tp0
]
q_values_tp0_min = torch.cat(q_values_tp0, dim=-1).min(dim=-1)[0]
policy_loss = torch.mean(log_pi_tp0 - q_values_tp0_min)
# critic loss (kl-divergence between categorical distributions)
actions_tp1, log_pi_tp1 = self.actor(
states_tp1, with_log_pi=True
)
log_pi_tp1 = log_pi_tp1 / self.reward_scale
logits_t = [x(states_t, actions_t) for x in self.critics]
logits_tp1 = [x(states_tp1, actions_tp1) for x in self.target_critics]
probs_tp1 = [F.softmax(x, dim=-1) for x in logits_tp1]
q_values_tp1 = [
torch.sum(x * self.z, dim=-1).unsqueeze_(-1) for x in probs_tp1
]
probs_ids_tp1_min = torch.cat(q_values_tp1, dim=-1).argmin(dim=1)
logits_tp1 = torch.cat([x.unsqueeze(-1) for x in logits_tp1], dim=-1)
logits_tp1 = logits_tp1[range(len(logits_tp1)), :, probs_ids_tp1_min]
gamma = self.gamma**self.n_step
z_target_tp1 = (self.z[None, :] - log_pi_tp1[:, None]).detach()
atoms_target_t = rewards_t + (1 - done_t) * gamma * z_target_tp1
value_loss = [
categorical_loss(
x, logits_tp1, atoms_target_t, self.z, self.delta_z,
self.v_min, self.v_max
) for x in logits_t
]
return policy_loss, value_loss
def _categorical_loss(self, states_t, actions_t, rewards_t, states_tp1,
done_t):
critic = self.critic
target_critic = self.target_critic
gammas = (self._gammas**self._n_step)[None, :, None]
# 1 x num_heads x 1
done_t = done_t[:, None, :] # B x 1 x 1
rewards_t = rewards_t[:, None, :] # B x 1 x 1
actions_t = actions_t[:, None, None, :] # B x 1 x 1 x 1
indices_t = actions_t.repeat(1, self._num_heads, 1, self.num_atoms)
# B x num_heads x 1 x num_atoms
logits_t = critic(states_t).gather(-2, indices_t).squeeze(-2)
# B x num_heads x num_atoms
all_logits_tp1 = target_critic(states_tp1).detach()
# B x num_heads x num_actions x num_atoms
q_values_tp1 = torch.sum(torch.softmax(all_logits_tp1, dim=-1) *
self.z,
dim=-1)
actions_tp1 = torch.argmax(q_values_tp1, dim=-1, keepdim=True)
# B x num_heads x 1
indices_tp1 = \
actions_tp1.unsqueeze(-1).repeat(1, 1, 1, self.num_atoms)
# B x num_heads x 1 x num_atoms
logits_tp1 = all_logits_tp1.gather(-2, indices_tp1).squeeze(-2)
# B x num_heads x num_atoms
atoms_target_t = rewards_t + (1 - done_t) * gammas * self.z
value_loss = categorical_loss(logits_t.view(-1, self.num_atoms),
logits_tp1.view(-1, self.num_atoms),
atoms_target_t.view(-1, self.num_atoms),
self.z, self.delta_z, self.v_min,
self.v_max)
return value_loss
def _categorical_loss(self, states_t, actions_t, rewards_t, states_tp1,
done_t):
gamma = self.gamma**self.n_step
# actor loss
logits_tp0 = self.critic(states_t, self.actor(states_t))
probs_tp0 = F.softmax(logits_tp0, dim=-1)
q_values_tp0 = torch.sum(probs_tp0 * self.z, dim=-1)
policy_loss = -torch.mean(q_values_tp0)
# critic loss (kl-divergence between categorical distributions)
logits_t = self.critic(states_t, actions_t)
logits_tp1 = self.target_critic(
states_tp1, self.target_actor(states_tp1)).detach()
atoms_target_t = rewards_t + (1 - done_t) * gamma * self.z
value_loss = categorical_loss(logits_t, logits_tp1, atoms_target_t,
self.z, self.delta_z, self.v_min,
self.v_max)
return policy_loss, value_loss
def _categorical_loss(self, states_t, actions_t, rewards_t, states_tp1,
done_t):
gamma = self._gamma**self._n_step
# critic loss (kl-divergence between categorical distributions)
indices_t = actions_t.repeat(1, self.num_atoms).unsqueeze(1)
logits_t = self.critic(states_t).gather(1, indices_t).squeeze(1)
all_logits_tp1 = self.target_critic(states_tp1).detach()
q_values_tp1 = torch.sum(F.softmax(all_logits_tp1, dim=-1) * self.z,
dim=-1)
actions_tp1 = torch.argmax(q_values_tp1, dim=-1, keepdim=True)
indices_tp1 = actions_tp1.repeat(1, self.num_atoms).unsqueeze(1)
logits_tp1 = all_logits_tp1.gather(1, indices_tp1).squeeze(1)
atoms_target_t = rewards_t + (1 - done_t) * gamma * self.z
value_loss = categorical_loss(logits_t, logits_tp1, atoms_target_t,
self.z, self.delta_z, self.v_min,
self.v_max)
return value_loss
def _categorical_loss(self, states_t, actions_t, rewards_t, states_tp1,
done_t):
# actor loss
actions_tp0 = self.actor(states_t)
logits_tp0 = [x(states_t, actions_tp0) for x in self.critics]
probs_tp0 = [F.softmax(x, dim=-1) for x in logits_tp0]
q_values_tp0 = [
torch.sum(x * self.z, dim=-1).unsqueeze_(-1) for x in probs_tp0
]
q_values_tp0_min = torch.cat(q_values_tp0, dim=-1).min(dim=-1)[0]
policy_loss = -torch.mean(q_values_tp0_min)
# critic loss (kl-divergence between categorical distributions)
actions_tp1 = self.target_actor(states_tp1).detach()
actions_tp1 = self._add_noise_to_actions(actions_tp1)
logits_t = [x(states_t, actions_t) for x in self.critics]
logits_tp1 = [x(states_tp1, actions_tp1) for x in self.target_critics]
probs_tp1 = [F.softmax(x, dim=-1) for x in logits_tp1]
q_values_tp1 = [
torch.sum(x * self.z, dim=-1).unsqueeze_(-1) for x in probs_tp1
]
probs_ids_tp1_min = torch.cat(q_values_tp1, dim=-1).argmin(dim=1)
logits_tp1 = torch.cat([x.unsqueeze(-1) for x in logits_tp1], dim=-1)
logits_tp1 = \
logits_tp1[range(len(logits_tp1)), :, probs_ids_tp1_min].detach()
gamma = self.gamma**self.n_step
atoms_target_t = rewards_t + (1 - done_t) * gamma * self.z
value_loss = [
categorical_loss(x, logits_tp1, atoms_target_t, self.z,
self.delta_z, self.v_min, self.v_max)
for x in logits_t
]
return policy_loss, value_loss