def calculate_quantile_loss(self, state_embeddings, tau_hats,
current_sa_quantile_hats, actions, rewards,
next_states, dones, weights):
assert not tau_hats.requires_grad
with torch.no_grad():
# NOTE: Current and target quantiles share the same proposed
# fractions to reduce computations. (i.e. next_tau_hats = tau_hats)
# Calculate Q values of next states.
if self.double_q_learning:
# Sample the noise of online network to decorrelate between
# the action selection and the quantile calculation.
self.online_net.sample_noise()
next_q = self.online_net.calculate_q(states=next_states)
else:
next_state_embeddings =\
self.target_net.calculate_state_embeddings(next_states)
next_q = self.target_net.calculate_q(
state_embeddings=next_state_embeddings,
fraction_net=self.online_net.fraction_net)
# Calculate greedy actions.
next_actions = torch.argmax(next_q, dim=1, keepdim=True)
assert next_actions.shape == (self.batch_size, 1)
# Calculate features of next states.
if self.double_q_learning:
next_state_embeddings =\
self.target_net.calculate_state_embeddings(next_states)
# Calculate quantile values of next states and actions at tau_hats.
next_sa_quantile_hats = evaluate_quantile_at_action(
self.target_net.calculate_quantiles(
taus=tau_hats, state_embeddings=next_state_embeddings),
next_actions).transpose(1, 2)
assert next_sa_quantile_hats.shape == (self.batch_size, 1, self.N)
# Calculate target quantile values.
target_sa_quantile_hats = rewards[..., None] + (
1.0 - dones[..., None]) * self.gamma_n * next_sa_quantile_hats
assert target_sa_quantile_hats.shape == (self.batch_size, 1,
self.N)
c_shape = current_sa_quantile_hats.shape
# print("!!!", c_shape)
# current_sa_quantile_hats = current_sa_quantile_hats[:, :, 0]#current_sa_quantile_hats.reshape(c_shape[0], c_shape[2], c_shape[1])
# print("calc_quan_loss", target_sa_quantile_hats.shape, current_sa_quantile_hats.shape)
td_errors = target_sa_quantile_hats - current_sa_quantile_hats
assert td_errors.shape == (self.batch_size, self.N, self.N)
quantile_huber_loss = calculate_quantile_huber_loss(
td_errors, tau_hats, weights, self.kappa)
return quantile_huber_loss, next_q.detach().mean().item(), \
td_errors.detach().abs()
def calculate_loss(self, states, actions, rewards, next_states, dones,
weights):
# Calculate quantile values of current states and actions at taus.
current_sa_quantiles = evaluate_quantile_at_action(
self.online_net(states=states),
actions)
assert current_sa_quantiles.shape == (self.batch_size, self.N, 1)
with torch.no_grad():
# Calculate Q values of next states.
if self.double_q_learning:
# Sample the noise of online network to decorrelate between
# the action selection and the quantile calculation.
self.online_net.sample_noise()
next_q = self.online_net.calculate_q(states=next_states)
else:
next_q = self.target_net.calculate_q(states=next_states)
# Calculate greedy actions.
next_actions = torch.argmax(next_q, dim=1, keepdim=True)
assert next_actions.shape == (self.batch_size, 1)
# Calculate quantile values of next states and actions at tau_hats.
next_sa_quantiles = evaluate_quantile_at_action(
self.target_net(states=next_states),
next_actions).transpose(1, 2)
assert next_sa_quantiles.shape == (self.batch_size, 1, self.N)
# Calculate target quantile values.
target_sa_quantiles = rewards[..., None] + (
1.0 - dones[..., None]) * self.gamma_n * next_sa_quantiles
assert target_sa_quantiles.shape == (self.batch_size, 1, self.N)
td_errors = target_sa_quantiles - current_sa_quantiles
assert td_errors.shape == (self.batch_size, self.N, self.N)
quantile_huber_loss = calculate_quantile_huber_loss(
td_errors, self.tau_hats, weights, self.kappa)
return quantile_huber_loss, next_q.detach().mean().item(), \
td_errors.detach().abs()