def update_target_networks(self):
ptu.soft_update_from_to(
self.qf1, self.target_qf1, self.soft_target_tau
)
ptu.soft_update_from_to(
self.qf2, self.target_qf2, self.soft_target_tau
)
def _update_target_networks(self):
if self.use_soft_update:
ptu.soft_update_from_to(self.policy, self.target_policy, self.tau)
ptu.soft_update_from_to(self.gcm, self.target_gcm, self.tau)
else:
if self._n_env_steps_total % self.target_hard_update_period == 0:
ptu.copy_model_params_from_to(self.gcm, self.target_gcm)
ptu.copy_model_params_from_to(self.policy, self.target_policy)
def _do_training(self):
batch = self.get_batch()
"""
Optimize Critic/Actor.
"""
rewards = batch['rewards']
terminals = batch['terminals']
obs = batch['observations']
actions = batch['actions']
next_obs = batch['next_observations']
_, _, v_pred = self.target_policy(next_obs, None)
y_target = self.reward_scale * rewards + (
1. - terminals) * self.discount * v_pred
y_target = y_target.detach()
mu, y_pred, v = self.policy(obs, actions)
policy_loss = self.policy_criterion(y_pred, y_target)
self.policy_optimizer.zero_grad()
policy_loss.backward()
self.policy_optimizer.step()
"""
Update Target Networks
"""
if self.use_soft_update:
ptu.soft_update_from_to(self.policy, self.target_policy, self.tau)
else:
if self._n_train_steps_total % self.target_hard_update_period == 0:
ptu.copy_model_params_from_to(self.policy, self.target_policy)
if self.need_to_update_eval_statistics:
self.need_to_update_eval_statistics = False
self.eval_statistics['Policy Loss'] = np.mean(
ptu.get_numpy(policy_loss))
self.eval_statistics.update(
create_stats_ordered_dict(
'Policy v',
ptu.get_numpy(v),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Policy mu',
ptu.get_numpy(mu),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Y targets',
ptu.get_numpy(y_target),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Y predictions',
ptu.get_numpy(y_pred),
))
def _do_training(self, n_steps_total):
raw_subtraj_batch, start_indices = (
self.replay_buffer.train_replay_buffer.random_subtrajectories(
self.num_subtrajs_per_batch
)
)
subtraj_batch = create_torch_subtraj_batch(raw_subtraj_batch)
if self.save_memory_gradients:
subtraj_batch['memories'].requires_grad = True
self.train_critic(subtraj_batch)
self.train_policy(subtraj_batch, start_indices)
if self.use_soft_update:
ptu.soft_update_from_to(self.policy, self.target_policy, self.tau)
ptu.soft_update_from_to(self.qf, self.target_qf, self.tau)
else:
if n_steps_total % self.target_hard_update_period == 0:
ptu.copy_model_params_from_to(self.qf, self.target_qf)
ptu.copy_model_params_from_to(self.policy, self.target_policy)
def train_from_torch(self, batch):
rewards = batch['rewards']
terminals = batch['terminals']
obs = batch['observations']
actions = batch['actions']
next_obs = batch['next_observations']
"""
Compute loss
"""
best_action_idxs = self.qf(next_obs).max(1, keepdim=True)[1]
target_q_values = self.target_qf(next_obs).gather(
1, best_action_idxs).detach()
y_target = rewards + (1. - terminals) * self.discount * target_q_values
y_target = y_target.detach()
# actions is a one-hot vector
y_pred = torch.sum(self.qf(obs) * actions, dim=1, keepdim=True)
qf_loss = self.qf_criterion(y_pred, y_target)
"""
Update networks
"""
self.qf_optimizer.zero_grad()
qf_loss.backward()
self.qf_optimizer.step()
"""
Soft target network updates
"""
if self._n_train_steps_total % self.target_update_period == 0:
ptu.soft_update_from_to(self.qf, self.target_qf,
self.soft_target_tau)
"""
Save some statistics for eval using just one batch.
"""
if self._need_to_update_eval_statistics:
self._need_to_update_eval_statistics = False
self.eval_statistics['QF Loss'] = np.mean(ptu.get_numpy(qf_loss))
self.eval_statistics.update(
create_stats_ordered_dict(
'Y Predictions',
ptu.get_numpy(y_pred),
))
def _do_training(self):
batch = self.get_batch()
rewards = batch['rewards']
terminals = batch['terminals']
obs = batch['observations']
actions = batch['actions']
next_obs = batch['next_observations']
goals = batch['goals']
num_steps_left = batch['num_steps_left']
q1_pred = self.qf1(
observations=obs,
actions=actions,
goals=goals,
num_steps_left=num_steps_left,
)
q2_pred = self.qf2(
observations=obs,
actions=actions,
goals=goals,
num_steps_left=num_steps_left,
)
# Make sure policy accounts for squashing functions like tanh correctly!
policy_outputs = self.policy(obs,
goals,
num_steps_left,
reparameterize=self.train_policy_with_reparameterization,
return_log_prob=True)
new_actions, policy_mean, policy_log_std, log_pi = policy_outputs[:4]
if not self.dense_rewards and not self.dense_log_pi:
log_pi = log_pi * terminals
"""
QF Loss
"""
target_v_values = self.target_vf(
observations=next_obs,
goals=goals,
num_steps_left=num_steps_left-1,
)
q_target = self.reward_scale * rewards + (1. - terminals) * self.discount * target_v_values
q_target = q_target.detach()
bellman_errors_1 = (q1_pred - q_target) ** 2
bellman_errors_2 = (q2_pred - q_target) ** 2
qf1_loss = bellman_errors_1.mean()
qf2_loss = bellman_errors_2.mean()
if self.use_automatic_entropy_tuning:
"""
Alpha Loss
"""
alpha_loss = -(self.log_alpha * (log_pi + self.target_entropy).detach()).mean()
self.alpha_optimizer.zero_grad()
alpha_loss.backward()
self.alpha_optimizer.step()
alpha = self.log_alpha.exp()
else:
alpha = 1
"""
VF Loss
"""
q1_new_actions = self.qf1(
observations=obs,
actions=new_actions,
goals=goals,
num_steps_left=num_steps_left,
)
q2_new_actions = self.qf2(
observations=obs,
actions=new_actions,
goals=goals,
num_steps_left=num_steps_left,
)
q_new_actions = torch.min(q1_new_actions, q2_new_actions)
v_target = q_new_actions - alpha * log_pi
v_pred = self.vf(
observations=obs,
goals=goals,
num_steps_left=num_steps_left,
)
v_target = v_target.detach()
bellman_errors = (v_pred - v_target) ** 2
vf_loss = bellman_errors.mean()
"""
Update networks
"""
self.qf1_optimizer.zero_grad()
qf1_loss.backward()
self.qf1_optimizer.step()
self.qf2_optimizer.zero_grad()
qf2_loss.backward()
self.qf2_optimizer.step()
self.vf_optimizer.zero_grad()
vf_loss.backward()
self.vf_optimizer.step()
"""
Policy Loss
"""
# paper says to do + but apparently that's a typo. Do Q - V.
if self.train_policy_with_reparameterization:
policy_loss = (alpha * log_pi - q_new_actions).mean()
else:
log_policy_target = q_new_actions - v_pred
policy_loss = (
log_pi * (alpha * log_pi - log_policy_target).detach()
).mean()
mean_reg_loss = self.policy_mean_reg_weight * (policy_mean ** 2).mean()
std_reg_loss = self.policy_std_reg_weight * (policy_log_std ** 2).mean()
pre_tanh_value = policy_outputs[-1]
pre_activation_reg_loss = self.policy_pre_activation_weight * (
(pre_tanh_value ** 2).sum(dim=1).mean()
)
policy_reg_loss = mean_reg_loss + std_reg_loss + pre_activation_reg_loss
policy_loss = policy_loss + policy_reg_loss
if self._n_train_steps_total % self.policy_update_period == 0:
self.policy_optimizer.zero_grad()
policy_loss.backward()
self.policy_optimizer.step()
if self._n_train_steps_total % self.target_update_period == 0:
ptu.soft_update_from_to(
self.vf, self.target_vf, self.soft_target_tau
)
"""
Save some statistics for eval
"""
if self.need_to_update_eval_statistics:
self.need_to_update_eval_statistics = False
"""
Eval should set this to None.
This way, these statistics are only computed for one batch.
"""
self.eval_statistics['QF1 Loss'] = np.mean(ptu.get_numpy(qf1_loss))
self.eval_statistics['QF2 Loss'] = np.mean(ptu.get_numpy(qf2_loss))
self.eval_statistics['VF Loss'] = np.mean(ptu.get_numpy(vf_loss))
self.eval_statistics['Policy Loss'] = np.mean(ptu.get_numpy(
policy_loss
))
self.eval_statistics.update(create_stats_ordered_dict(
'Q1 Predictions',
ptu.get_numpy(q1_pred),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Q2 Predictions',
ptu.get_numpy(q2_pred),
))
self.eval_statistics.update(create_stats_ordered_dict(
'V Predictions',
ptu.get_numpy(v_pred),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Log Pis',
ptu.get_numpy(log_pi),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Policy mu',
ptu.get_numpy(policy_mean),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Policy log std',
ptu.get_numpy(policy_log_std),
))
if self.use_automatic_entropy_tuning:
self.eval_statistics['Alpha'] = ptu.get_numpy(alpha)[0]
self.eval_statistics['Alpha Loss'] = ptu.get_numpy(alpha_loss)[0]
def _do_training(self):
batch = self.get_batch()
rewards = batch['rewards']
terminals = batch['terminals']
obs = batch['observations']
actions = batch['actions']
next_obs = batch['next_observations']
goals = batch['goals']
num_steps_left = batch['num_steps_left']
"""
Critic operations.
"""
next_actions = self.target_policy(
observations=next_obs,
goals=goals,
num_steps_left=num_steps_left - 1,
)
noise = torch.normal(
torch.zeros_like(next_actions),
self.target_policy_noise,
)
noise = torch.clamp(noise, -self.target_policy_noise_clip,
self.target_policy_noise_clip)
noisy_next_actions = next_actions + noise
target_q1_values = self.target_qf1(
observations=next_obs,
actions=noisy_next_actions,
goals=goals,
num_steps_left=num_steps_left - 1,
)
target_q2_values = self.target_qf2(
observations=next_obs,
actions=noisy_next_actions,
goals=goals,
num_steps_left=num_steps_left - 1,
)
target_q_values = torch.min(target_q1_values, target_q2_values)
q_target = self.reward_scale * rewards + (
1. - terminals) * self.discount * target_q_values
q_target = q_target.detach()
q1_pred = self.qf1(
observations=obs,
actions=actions,
goals=goals,
num_steps_left=num_steps_left,
)
q2_pred = self.qf2(
observations=obs,
actions=actions,
goals=goals,
num_steps_left=num_steps_left,
)
bellman_errors_1 = (q1_pred - q_target)**2
bellman_errors_2 = (q2_pred - q_target)**2
qf1_loss = bellman_errors_1.mean()
qf2_loss = bellman_errors_2.mean()
"""
Update Networks
"""
self.qf1_optimizer.zero_grad()
qf1_loss.backward()
self.qf1_optimizer.step()
self.qf2_optimizer.zero_grad()
qf2_loss.backward()
self.qf2_optimizer.step()
policy_actions, pre_tanh_value = self.policy(
obs,
goals,
num_steps_left,
return_preactivations=True,
)
policy_saturation_cost = F.relu(torch.abs(pre_tanh_value) - 20.0)
q_output = self.qf1(
observations=obs,
actions=policy_actions,
num_steps_left=num_steps_left,
goals=goals,
)
policy_loss = -q_output.mean()
if self.use_policy_saturation_cost:
policy_loss = policy_loss + policy_saturation_cost.mean()
if self._n_train_steps_total % self.policy_and_target_update_period == 0:
self.policy_optimizer.zero_grad()
policy_loss.backward()
self.policy_optimizer.step()
ptu.soft_update_from_to(self.policy, self.target_policy, self.tau)
ptu.soft_update_from_to(self.qf1, self.target_qf1, self.tau)
ptu.soft_update_from_to(self.qf2, self.target_qf2, self.tau)
if self.need_to_update_eval_statistics:
self.need_to_update_eval_statistics = False
self.eval_statistics = OrderedDict()
self.eval_statistics['QF1 Loss'] = np.mean(ptu.get_numpy(qf1_loss))
self.eval_statistics['QF2 Loss'] = np.mean(ptu.get_numpy(qf2_loss))
self.eval_statistics['Policy Loss'] = np.mean(
ptu.get_numpy(policy_loss))
self.eval_statistics.update(
create_stats_ordered_dict(
'Q1 Predictions',
ptu.get_numpy(q1_pred),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Q2 Predictions',
ptu.get_numpy(q2_pred),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Q Targets',
ptu.get_numpy(q_target),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Bellman1 Errors',
ptu.get_numpy(bellman_errors_1),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Bellman2 Errors',
ptu.get_numpy(bellman_errors_2),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Policy Saturation Cost',
ptu.get_numpy(policy_saturation_cost),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Policy Action',
ptu.get_numpy(policy_actions),
exclude_abs=False,
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Policy Action Pre-tanh',
ptu.get_numpy(pre_tanh_value),
exclude_abs=False,
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Policy Q Output',
ptu.get_numpy(q_output),
exclude_abs=False,
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Policy Noisy Next Actions',
ptu.get_numpy(noisy_next_actions),
exclude_abs=False,
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Replay Buffer Actions',
ptu.get_numpy(actions),
exclude_abs=False,
))
be_np = ptu.get_numpy(bellman_errors_1)
num_steps_left_np = ptu.get_numpy(num_steps_left)
### tau == 0 ###
idx_0 = np.argwhere(num_steps_left_np == 0)
be_0 = 0
if len(idx_0) > 0:
be_0 = be_np[idx_0[:, 0]]
self.eval_statistics['QF1 Loss tau=0'] = np.mean(be_0)
### tau == 1 ###
idx_1 = np.argwhere(num_steps_left_np == 1)
be_1 = 0
if len(idx_1) > 0:
be_1 = be_np[idx_1[:, 0]]
self.eval_statistics['QF1 Loss tau=1'] = np.mean(be_1)
### tau in [2, 5) ###
idx_2_to_5 = np.argwhere(
np.logical_and(num_steps_left_np >= 2, num_steps_left_np < 5))
be_2_to_5 = 0
if len(idx_2_to_5) > 0:
be_2_to_5 = be_np[idx_2_to_5[:, 0]]
self.eval_statistics['QF1 Loss tau=2_to_5'] = np.mean(be_2_to_5)
### tau in [5, 10) ###
idx_5_to_10 = np.argwhere(
np.logical_and(num_steps_left_np >= 5, num_steps_left_np < 10))
be_5_to_10 = 0
if len(idx_5_to_10) > 0:
be_5_to_10 = be_np[idx_5_to_10[:, 0]]
self.eval_statistics['QF1 Loss tau=5_to_10'] = np.mean(be_5_to_10)
### tau in [10, max_tau] ###
idx_10_to_end = np.argwhere(
np.logical_and(num_steps_left_np >= 10,
num_steps_left_np < self.max_tau + 1))
be_10_to_end = 0
if len(idx_10_to_end) > 0:
be_10_to_end = be_np[idx_10_to_end[:, 0]]
self.eval_statistics['QF1 Loss tau=10_to_end'] = np.mean(
be_10_to_end)
def train_from_torch(self, batch):
rewards = batch['rewards']
terminals = batch['terminals']
obs = batch['observations']
actions = batch['actions']
next_obs = batch['next_observations']
"""
Classifier and Policy
"""
class_actions = self.policy(obs)
class_prob = self.classifier(obs, actions)
prob_target = 1 + rewards[:, -1]
neg_log_prob = - torch.log(self.classifier(obs, class_actions))
policy_loss = (neg_log_prob).mean()
"""
QF Loss
"""
q1_pred = self.qf1(obs, actions)
q2_pred = self.qf2(obs, actions)
# Make sure policy accounts for squashing functions like tanh correctly!
new_next_actions, _, _, new_log_pi, *_ = self.policy(
next_obs, reparameterize=True, return_log_prob=True,
)
target_q_values = torch.min(
self.target_qf1(next_obs, new_next_actions),
self.target_qf2(next_obs, new_next_actions),
) - alpha * new_log_pi
q_target = self.reward_scale * rewards + (1. - terminals) * self.discount * target_q_values
qf1_loss = self.qf_criterion(q1_pred, q_target.detach())
qf2_loss = self.qf_criterion(q2_pred, q_target.detach())
"""
Update networks
"""
self.qf1_optimizer.zero_grad()
qf1_loss.backward()
self.qf1_optimizer.step()
self.qf2_optimizer.zero_grad()
qf2_loss.backward()
self.qf2_optimizer.step()
self.policy_optimizer.zero_grad()
policy_loss.backward()
self.policy_optimizer.step()
"""
Soft Updates
"""
if self._n_train_steps_total % self.target_update_period == 0:
ptu.soft_update_from_to(
self.qf1, self.target_qf1, self.soft_target_tau
)
ptu.soft_update_from_to(
self.qf2, self.target_qf2, self.soft_target_tau
)
"""
Save some statistics for eval
"""
if self._need_to_update_eval_statistics:
self._need_to_update_eval_statistics = False
"""
Eval should set this to None.
This way, these statistics are only computed for one batch.
"""
policy_loss = (log_pi - q_new_actions).mean()
self.eval_statistics['QF1 Loss'] = np.mean(ptu.get_numpy(qf1_loss))
self.eval_statistics['QF2 Loss'] = np.mean(ptu.get_numpy(qf2_loss))
self.eval_statistics['Policy Loss'] = np.mean(ptu.get_numpy(
policy_loss
))
self.eval_statistics.update(create_stats_ordered_dict(
'Q1 Predictions',
ptu.get_numpy(q1_pred),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Q2 Predictions',
ptu.get_numpy(q2_pred),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Q Targets',
ptu.get_numpy(q_target),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Log Pis',
ptu.get_numpy(log_pi),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Policy mu',
ptu.get_numpy(policy_mean),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Policy log std',
ptu.get_numpy(policy_log_std),
))
if self.use_automatic_entropy_tuning:
self.eval_statistics['Alpha'] = alpha.item()
self.eval_statistics['Alpha Loss'] = alpha_loss.item()
self._n_train_steps_total += 1
def _train_given_data(
self,
rewards,
terminals,
obs,
actions,
next_obs,
logger_prefix="",
):
"""
Critic operations.
"""
next_actions = self.target_policy(next_obs)
noise = torch.normal(
torch.zeros_like(next_actions),
self.target_policy_noise,
)
noise = torch.clamp(noise, -self.target_policy_noise_clip,
self.target_policy_noise_clip)
noisy_next_actions = next_actions + noise
target_q1_values = self.target_qf1(next_obs, noisy_next_actions)
target_q2_values = self.target_qf2(next_obs, noisy_next_actions)
target_q_values = torch.min(target_q1_values, target_q2_values)
q_target = self.reward_scale * rewards + (
1. - terminals) * self.discount * target_q_values
q_target = q_target.detach()
q1_pred = self.qf1(obs, actions)
bellman_errors_1 = (q1_pred - q_target)**2
qf1_loss = bellman_errors_1.mean()
q2_pred = self.qf2(obs, actions)
bellman_errors_2 = (q2_pred - q_target)**2
qf2_loss = bellman_errors_2.mean()
"""
Update Networks
"""
self.qf1_optimizer.zero_grad()
qf1_loss.backward()
self.qf1_optimizer.step()
self.qf2_optimizer.zero_grad()
qf2_loss.backward()
self.qf2_optimizer.step()
policy_actions = policy_loss = None
if self._n_train_steps_total % self.policy_and_target_update_period == 0:
policy_actions = self.policy(obs)
q_output = self.qf1(obs, policy_actions)
policy_loss = -q_output.mean()
self.policy_optimizer.zero_grad()
policy_loss.backward()
self.policy_optimizer.step()
ptu.soft_update_from_to(self.policy, self.target_policy, self.tau)
ptu.soft_update_from_to(self.qf1, self.target_qf1, self.tau)
ptu.soft_update_from_to(self.qf2, self.target_qf2, self.tau)
if self.need_to_update_eval_statistics:
self.need_to_update_eval_statistics = False
if policy_loss is None:
policy_actions = self.policy(obs)
q_output = self.qf1(obs, policy_actions)
policy_loss = -q_output.mean()
self.eval_statistics[logger_prefix + 'QF1 Loss'] = np.mean(
ptu.get_numpy(qf1_loss))
self.eval_statistics[logger_prefix + 'QF2 Loss'] = np.mean(
ptu.get_numpy(qf2_loss))
self.eval_statistics[logger_prefix + 'Policy Loss'] = np.mean(
ptu.get_numpy(policy_loss))
self.eval_statistics.update(
create_stats_ordered_dict(
logger_prefix + 'Q1 Predictions',
ptu.get_numpy(q1_pred),
))
self.eval_statistics.update(
create_stats_ordered_dict(
logger_prefix + 'Q2 Predictions',
ptu.get_numpy(q2_pred),
))
self.eval_statistics.update(
create_stats_ordered_dict(
logger_prefix + 'Q Targets',
ptu.get_numpy(q_target),
))
self.eval_statistics.update(
create_stats_ordered_dict(
logger_prefix + 'Bellman Errors 1',
ptu.get_numpy(bellman_errors_1),
))
self.eval_statistics.update(
create_stats_ordered_dict(
logger_prefix + 'Bellman Errors 2',
ptu.get_numpy(bellman_errors_2),
))
self.eval_statistics.update(
create_stats_ordered_dict(
logger_prefix + 'Policy Action',
ptu.get_numpy(policy_actions),
))
def _do_training(self):
batch = self.get_batch()
rewards = batch['rewards']
terminals = batch['terminals']
obs = batch['observations']
actions = batch['actions']
next_obs = batch['next_observations']
goals = batch['goals']
num_steps_left = batch['num_steps_left']
"""
Critic operations.
"""
next_actions = self.target_policy(
observations=next_obs,
goals=goals,
num_steps_left=num_steps_left - 1,
)
noise = torch.normal(
torch.zeros_like(next_actions),
self.target_policy_noise,
)
noise = torch.clamp(noise, -self.target_policy_noise_clip,
self.target_policy_noise_clip)
noisy_next_actions = next_actions + noise
target_q1_values = self.target_qf1(
observations=next_obs,
actions=noisy_next_actions,
goals=goals,
num_steps_left=num_steps_left - 1,
)
target_q2_values = self.target_qf2(
observations=next_obs,
actions=noisy_next_actions,
goals=goals,
num_steps_left=num_steps_left - 1,
)
target_q_values = torch.min(target_q1_values, target_q2_values)
q_target = self.reward_scale * rewards + (
1. - terminals) * self.discount * target_q_values
q_target = q_target.detach()
q1_pred = self.qf1(
observations=obs,
actions=actions,
goals=goals,
num_steps_left=num_steps_left,
)
q2_pred = self.qf2(
observations=obs,
actions=actions,
goals=goals,
num_steps_left=num_steps_left,
)
bellman_errors_1 = (q1_pred - q_target)**2
bellman_errors_2 = (q2_pred - q_target)**2
qf1_loss = bellman_errors_1.mean()
qf2_loss = bellman_errors_2.mean()
"""
Update Networks
"""
self.qf1_optimizer.zero_grad()
qf1_loss.backward()
self.qf1_optimizer.step()
self.qf2_optimizer.zero_grad()
qf2_loss.backward()
self.qf2_optimizer.step()
policy_actions, pre_tanh_value = self.policy(
obs,
goals,
num_steps_left,
return_preactivations=True,
)
q_output = self.qf1(
observations=obs,
actions=policy_actions,
num_steps_left=num_steps_left,
goals=goals,
)
policy_loss = -q_output.mean()
if self._n_train_steps_total % self.policy_and_target_update_period == 0:
self.policy_optimizer.zero_grad()
policy_loss.backward()
self.policy_optimizer.step()
ptu.soft_update_from_to(self.policy, self.target_policy, self.tau)
ptu.soft_update_from_to(self.qf1, self.target_qf1, self.tau)
ptu.soft_update_from_to(self.qf2, self.target_qf2, self.tau)
if self.need_to_update_eval_statistics:
self.need_to_update_eval_statistics = False
self.eval_statistics = OrderedDict()
self.eval_statistics['QF1 Loss'] = np.mean(ptu.get_numpy(qf1_loss))
self.eval_statistics['QF2 Loss'] = np.mean(ptu.get_numpy(qf2_loss))
self.eval_statistics['Policy Loss'] = np.mean(
ptu.get_numpy(policy_loss))
self.eval_statistics.update(
create_stats_ordered_dict(
'Q1 Predictions',
ptu.get_numpy(q1_pred),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Q2 Predictions',
ptu.get_numpy(q2_pred),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Q Targets',
ptu.get_numpy(q_target),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Bellman1 Errors',
ptu.get_numpy(bellman_errors_1),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Bellman2 Errors',
ptu.get_numpy(bellman_errors_2),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Policy Action',
ptu.get_numpy(policy_actions),
))
def train_from_torch(self, batch):
self._current_epoch += 1
rewards = batch['rewards']
terminals = batch['terminals']
obs = batch['observations']
obs = obs.reshape((obs.shape[0], ) + (3, 48, 48))
if not self.discrete:
actions = batch['actions']
else:
actions = batch['actions'].argmax(dim=-1)
next_obs = batch['next_observations']
next_obs = next_obs.reshape((next_obs.shape[0], ) + (3, 48, 48))
"""
Policy and Alpha Loss
"""
if self.discrete:
new_obs_actions, pi_probs, log_pi, entropies = self.policy(
obs, None, return_log_prob=True)
new_next_actions, pi_next_probs, new_log_pi, next_entropies = self.policy(
next_obs, None, return_log_prob=True)
q_vector = self.qf1.q_vector(obs)
q2_vector = self.qf2.q_vector(obs)
q_next_vector = self.qf1.q_vector(next_obs)
q2_next_vector = self.qf2.q_vector(next_obs)
else:
new_obs_actions, policy_mean, policy_log_std, log_pi, *_ = self.policy(
obs,
None,
reparameterize=True,
return_log_prob=True,
)
if self.use_automatic_entropy_tuning:
alpha_loss = -(self.log_alpha *
(log_pi + self.target_entropy).detach()).mean()
self.alpha_optimizer.zero_grad()
alpha_loss.backward()
self.alpha_optimizer.step()
alpha = self.log_alpha.exp()
else:
alpha_loss = 0
alpha = 1
if not self.discrete:
if self.num_qs == 1:
q_new_actions = self.qf1(obs, None, new_obs_actions)
else:
q_new_actions = torch.min(
self.qf1(obs, None, new_obs_actions),
self.qf2(obs, None, new_obs_actions),
)
if self.discrete:
target_q_values = torch.min(q_vector, q2_vector)
policy_loss = -((target_q_values * pi_probs).sum(dim=-1) +
alpha * entropies).mean()
else:
policy_loss = (alpha * log_pi - q_new_actions).mean()
if self._current_epoch < self.policy_eval_start:
"""Start with BC"""
policy_log_prob = self.policy.log_prob(obs, None, actions)
policy_loss = (alpha * log_pi - policy_log_prob).mean()
# print ('Policy Loss: ', policy_loss.item())
"""
QF Loss
"""
q1_pred = self.qf1(obs, None, actions)
if self.num_qs > 1:
q2_pred = self.qf2(obs, None, actions)
# Make sure policy accounts for squashing functions like tanh correctly!
if not self.discrete:
new_next_actions, _, _, new_log_pi, *_ = self.policy(
next_obs,
None,
reparameterize=True,
return_log_prob=True,
)
new_curr_actions, _, _, new_curr_log_pi, *_ = self.policy(
obs,
None,
reparameterize=True,
return_log_prob=True,
)
else:
new_curr_actions, pi_curr_probs, new_curr_log_pi, new_curr_entropies = self.policy(
obs, None, return_log_prob=True)
if not self.max_q_backup:
if not self.discrete:
if self.num_qs == 1:
target_q_values = self.target_qf1(next_obs, None,
new_next_actions)
else:
target_q_values = torch.min(
self.target_qf1(next_obs, None, new_next_actions),
self.target_qf2(next_obs, None, new_next_actions),
)
else:
target_q_values = torch.min(
(self.target_qf1.q_vector(next_obs) *
pi_next_probs).sum(dim=-1),
(self.target_qf2.q_vector(next_obs) *
pi_next_probs).sum(dim=-1))
target_q_values = target_q_values.unsqueeze(-1)
if not self.deterministic_backup:
target_q_values = target_q_values - alpha * new_log_pi
if self.max_q_backup:
"""when using max q backup"""
if not self.discrete:
next_actions_temp, _ = self._get_policy_actions(
next_obs, num_actions=10, network=self.policy)
target_qf1_values = \
self._get_tensor_values(next_obs, next_actions_temp,
network=self.target_qf1).max(1)[0].view(
-1, 1)
target_qf2_values = \
self._get_tensor_values(next_obs, next_actions_temp,
network=self.target_qf2).max(1)[0].view(
-1, 1)
target_q_values = torch.min(
target_qf1_values, target_qf2_values
) # + torch.max(target_qf1_values, target_qf2_values) * 0.25
else:
target_qf1_values = \
self.target_qf1.q_vector(next_obs).max(dim=-1)[0]
target_qf2_values = \
self.target_qf2.q_vector(next_obs).max(dim=-1)[0]
target_q_values = torch.min(target_qf1_values,
target_qf2_values).unsqueeze(-1)
q_target = self.reward_scale * rewards + (
1. - terminals) * self.discount * target_q_values
# Only detach if we are not using Bellman residual and not otherwise
if self._use_target_nets:
q_target = q_target.detach()
qf1_loss = self.qf_criterion(q1_pred, q_target)
if self.num_qs > 1:
qf2_loss = self.qf_criterion(q2_pred, q_target)
if self.hinge_bellman:
qf1_loss = self.softplus(q_target - q1_pred).mean()
qf2_loss = self.softplus(q_target - q2_pred).mean()
## add min_q
if self.with_min_q:
if not self.discrete:
random_actions_tensor = torch.FloatTensor(
q2_pred.shape[0] * self.num_random,
actions.shape[-1]).uniform_(-1, 1).cuda()
curr_actions_tensor, curr_log_pis = self._get_policy_actions(
obs, num_actions=self.num_random, network=self.policy)
new_curr_actions_tensor, new_log_pis = self._get_policy_actions(
next_obs, num_actions=self.num_random, network=self.policy)
q1_rand = self._get_tensor_values(obs,
random_actions_tensor,
network=self.qf1)
q2_rand = self._get_tensor_values(obs,
random_actions_tensor,
network=self.qf2)
q1_curr_actions = self._get_tensor_values(obs,
curr_actions_tensor,
network=self.qf1)
q2_curr_actions = self._get_tensor_values(obs,
curr_actions_tensor,
network=self.qf2)
q1_next_actions = self._get_tensor_values(
obs, new_curr_actions_tensor, network=self.qf1)
q2_next_actions = self._get_tensor_values(
obs, new_curr_actions_tensor, network=self.qf2)
# q1_next_states_actions = self._get_tensor_values(next_obs, new_curr_actions_tensor, network=self.qf1)
# q2_next_states_actions = self._get_tensor_values(next_obs, new_curr_actions_tensor, network=self.qf2)
cat_q1 = torch.cat([
q1_rand,
q1_pred.unsqueeze(1), q1_next_actions, q1_curr_actions
], 1)
cat_q2 = torch.cat([
q2_rand,
q2_pred.unsqueeze(1), q2_next_actions, q2_curr_actions
], 1)
std_q1 = torch.std(cat_q1, dim=1)
std_q2 = torch.std(cat_q2, dim=1)
if self.min_q_version == 3:
# importance sammpled version
random_density = np.log(0.5**curr_actions_tensor.shape[-1])
cat_q1 = torch.cat([
q1_rand - random_density,
q1_next_actions - new_log_pis.detach(),
q1_curr_actions - curr_log_pis.detach()
], 1)
cat_q2 = torch.cat([
q2_rand - random_density,
q2_next_actions - new_log_pis.detach(),
q2_curr_actions - curr_log_pis.detach()
], 1)
if self.min_q_version == 0:
min_qf1_loss = cat_q1.mean() * self.min_q_weight
min_qf2_loss = cat_q2.mean() * self.min_q_weight
elif self.min_q_version == 1:
"""Expectation under softmax distribution"""
softmax_dist_1 = self.softmax(
cat_q1 / self.temp).detach() * self.temp
softmax_dist_2 = self.softmax(
cat_q2 / self.temp).detach() * self.temp
min_qf1_loss = (cat_q1 *
softmax_dist_1).mean() * self.min_q_weight
min_qf2_loss = (cat_q2 *
softmax_dist_2).mean() * self.min_q_weight
elif self.min_q_version == 2 or self.min_q_version == 3:
"""log sum exp for the min"""
min_qf1_loss = torch.logsumexp(
cat_q1 / self.temp,
dim=1,
).mean() * self.min_q_weight * self.temp
min_qf2_loss = torch.logsumexp(
cat_q2 / self.temp,
dim=1,
).mean() * self.min_q_weight * self.temp
if self.data_subtract:
"""Subtract the log likelihood of data"""
min_qf1_loss = min_qf1_loss - q1_pred.mean(
) * self.min_q_weight
min_qf2_loss = min_qf2_loss - q2_pred.mean(
) * self.min_q_weight
else:
q1_policy = (q_vector * pi_probs).sum(dim=-1)
q2_policy = (q2_vector * pi_probs).sum(dim=-1)
q1_next_actions = (q_next_vector * pi_next_probs).sum(dim=-1)
q2_next_actions = (q2_next_vector * pi_next_probs).sum(dim=-1)
if self.min_q_version == 0:
min_qf1_loss = (q1_policy.mean() + q1_next_actions.mean() +
q_vector.mean() + q_next_vector.mean()
).mean() * self.min_q_weight
min_qf2_loss = (q2_policy.mean() + q1_next_actions.mean() +
q2_vector.mean() + q2_next_vector.mean()
).mean() * self.min_q_weight
elif self.min_q_version == 1:
min_qf1_loss = (q_vector.mean() + q_next_vector.mean()
).mean() * self.min_q_weight
min_qf2_loss = (q2_vector.mean() + q2_next_vector.mean()
).mean() * self.min_q_weight
else:
softmax_dist_q1 = self.softmax(
q_vector / self.temp).detach() * self.temp
softmax_dist_q2 = self.softmax(
q2_vector / self.temp).detach() * self.temp
min_qf1_loss = (q_vector *
softmax_dist_q1).mean() * self.min_q_weight
min_qf2_loss = (q2_vector *
softmax_dist_q2).mean() * self.min_q_weight
if self.data_subtract:
min_qf1_loss = min_qf1_loss - q1_pred.mean(
) * self.min_q_weight
min_qf2_loss = min_qf2_loss - q2_pred.mean(
) * self.min_q_weight
std_q1 = torch.std(q_vector, dim=-1)
std_q2 = torch.std(q2_vector, dim=-1)
q1_on_policy = q1_policy.mean()
q2_on_policy = q2_policy.mean()
q1_random = q_vector.mean()
q2_random = q2_vector.mean()
q1_next_actions_mean = q1_next_actions.mean()
q2_next_actions_mean = q2_next_actions.mean()
if self.use_projected_grad:
min_qf1_grad = torch.autograd.grad(
min_qf1_loss,
inputs=[p for p in self.qf1.parameters()],
create_graph=True,
retain_graph=True,
only_inputs=True)
min_qf2_grad = torch.autograd.grad(
min_qf2_loss,
inputs=[p for p in self.qf2.parameters()],
create_graph=True,
retain_graph=True,
only_inputs=True)
qf1_loss_grad = torch.autograd.grad(
qf1_loss,
inputs=[p for p in self.qf1.parameters()],
create_graph=True,
retain_graph=True,
only_inputs=True)
qf2_loss_grad = torch.autograd.grad(
qf2_loss,
inputs=[p for p in self.qf2.parameters()],
create_graph=True,
retain_graph=True,
only_inputs=True)
# this is for the offline setting
# qf1_total_grad = self.compute_mt_grad(qf1_loss_grad, min_qf1_grad)
# qf2_total_grad = self.compute_mt_grad(qf2_loss_grad, min_qf2_grad)
qf1_total_grad = self.compute_new_grad(min_qf1_grad,
qf1_loss_grad)
qf2_total_grad = self.compute_new_grad(min_qf2_grad,
qf2_loss_grad)
else:
if self.with_lagrange:
alpha_prime = torch.clamp(self.log_alpha_prime.exp(),
min=0,
max=2000000.0)
orig_min_qf1_loss = min_qf1_loss
orig_min_qf2_loss = min_qf2_loss
min_qf1_loss = alpha_prime * (min_qf1_loss -
self.target_action_gap)
min_qf2_loss = alpha_prime * (min_qf2_loss -
self.target_action_gap)
self.alpha_prime_optimizer.zero_grad()
alpha_prime_loss = -0.5 * (min_qf1_loss + min_qf2_loss)
alpha_prime_loss.backward(retain_graph=True)
self.alpha_prime_optimizer.step()
qf1_loss = qf1_loss + min_qf1_loss
qf2_loss = qf2_loss + min_qf2_loss
"""
Update networks
"""
# Update the Q-functions iff
self._num_q_update_steps += 1
self.qf1_optimizer.zero_grad()
qf1_loss.backward(retain_graph=True)
if self.with_min_q and self.use_projected_grad:
for (p, proj_grad) in zip(self.qf1.parameters(), qf1_total_grad):
p.grad.data = proj_grad
self.qf1_optimizer.step()
if self.num_qs > 1:
self.qf2_optimizer.zero_grad()
qf2_loss.backward(retain_graph=True)
if self.with_min_q and self.use_projected_grad:
for (p, proj_grad) in zip(self.qf2.parameters(),
qf2_total_grad):
p.grad.data = proj_grad
self.qf2_optimizer.step()
self._num_policy_update_steps += 1
self.policy_optimizer.zero_grad()
policy_loss.backward(retain_graph=False)
self.policy_optimizer.step()
"""
Soft Updates
"""
if self._use_target_nets:
if self._n_train_steps_total % self.target_update_period == 0:
ptu.soft_update_from_to(self.qf1, self.target_qf1,
self.soft_target_tau)
if self.num_qs > 1:
ptu.soft_update_from_to(self.qf2, self.target_qf2,
self.soft_target_tau)
"""
Save some statistics for eval
"""
if self._need_to_update_eval_statistics:
self._need_to_update_eval_statistics = False
"""
Eval should set this to None.
This way, these statistics are only computed for one batch.
"""
if not self.discrete:
policy_loss = (log_pi - q_new_actions).mean()
else:
target_q_values = torch.min(q_vector, q2_vector)
policy_loss = -((target_q_values * pi_probs).sum(dim=-1) +
alpha * entropies).mean()
self.eval_statistics['QF1 Loss'] = np.mean(ptu.get_numpy(qf1_loss))
if self.num_qs > 1:
self.eval_statistics['QF2 Loss'] = np.mean(
ptu.get_numpy(qf2_loss))
if self.with_min_q and not self.discrete:
self.eval_statistics['Std QF1 values'] = np.mean(
ptu.get_numpy(std_q1))
self.eval_statistics['Std QF2 values'] = np.mean(
ptu.get_numpy(std_q2))
self.eval_statistics.update(
create_stats_ordered_dict(
'QF1 in-distribution values',
ptu.get_numpy(q1_curr_actions),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'QF2 in-distribution values',
ptu.get_numpy(q2_curr_actions),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'QF1 random values',
ptu.get_numpy(q1_rand),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'QF2 random values',
ptu.get_numpy(q2_rand),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'QF1 next_actions values',
ptu.get_numpy(q1_next_actions),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'QF2 next_actions values',
ptu.get_numpy(q2_next_actions),
))
elif self.with_min_q and self.discrete:
self.eval_statistics['Std QF1 values'] = np.mean(
ptu.get_numpy(std_q1))
self.eval_statistics['Std QF2 values'] = np.mean(
ptu.get_numpy(std_q2))
self.eval_statistics['QF1 on policy average'] = np.mean(
ptu.get_numpy(q1_on_policy))
self.eval_statistics['QF2 on policy average'] = np.mean(
ptu.get_numpy(q2_on_policy))
self.eval_statistics['QF1 random average'] = np.mean(
ptu.get_numpy(q1_random))
self.eval_statistics['QF2 random average'] = np.mean(
ptu.get_numpy(q2_random))
self.eval_statistics[
'QF1 next_actions_mean average'] = np.mean(
ptu.get_numpy(q1_next_actions_mean))
self.eval_statistics[
'QF2 next_actions_mean average'] = np.mean(
ptu.get_numpy(q2_next_actions_mean))
self.eval_statistics['Num Q Updates'] = self._num_q_update_steps
self.eval_statistics[
'Num Policy Updates'] = self._num_policy_update_steps
self.eval_statistics['Policy Loss'] = np.mean(
ptu.get_numpy(policy_loss))
self.eval_statistics.update(
create_stats_ordered_dict(
'Q1 Predictions',
ptu.get_numpy(q1_pred),
))
if self.num_qs > 1:
self.eval_statistics.update(
create_stats_ordered_dict(
'Q2 Predictions',
ptu.get_numpy(q2_pred),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Q Targets',
ptu.get_numpy(q_target),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Log Pis',
ptu.get_numpy(log_pi),
))
if not self.discrete:
self.eval_statistics.update(
create_stats_ordered_dict(
'Policy mu',
ptu.get_numpy(policy_mean),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Policy log std',
ptu.get_numpy(policy_log_std),
))
else:
self.eval_statistics['Policy entropy'] = ptu.get_numpy(
entropies).mean()
if self.use_automatic_entropy_tuning:
self.eval_statistics['Alpha'] = alpha.item()
self.eval_statistics['Alpha Loss'] = alpha_loss.item()
if self.with_lagrange:
self.eval_statistics['Alpha Prime'] = alpha_prime.item()
self.eval_statistics[
'Alpha Prime Loss'] = alpha_prime_loss.item()
self.eval_statistics['Min Q1 Loss'] = orig_min_qf1_loss.item()
self.eval_statistics['Min Q2 Loss'] = orig_min_qf2_loss.item()
self._n_train_steps_total += 1
def train_from_torch(self, batch):
rewards = batch['rewards']
terminals = batch['terminals']
obs = batch['observations']
actions = batch['actions']
next_obs = batch['next_observations']
weights = batch.get('weights', None)
if self.reward_transform:
rewards = self.reward_transform(rewards)
if self.terminal_transform:
terminals = self.terminal_transform(terminals)
"""
Policy and Alpha Loss
"""
new_obs_actions, policy_mean, policy_log_std, log_pi, entropy, policy_std, mean_action_log_prob, pretanh_value, dist = self.policy(
obs, reparameterize=True, return_log_prob=True,
)
if self.use_automatic_entropy_tuning:
alpha_loss = -(self.log_alpha * (log_pi + self.target_entropy).detach()).mean()
self.alpha_optimizer.zero_grad()
alpha_loss.backward()
self.alpha_optimizer.step()
alpha = self.log_alpha.exp()
else:
alpha_loss = 0
alpha = self.alpha
"""
QF Loss
"""
q1_pred = self.qf1(obs, actions)
q2_pred = self.qf2(obs, actions)
# Make sure policy accounts for squashing functions like tanh correctly!
new_next_actions, _, _, new_log_pi, *_ = self.policy(
next_obs, reparameterize=True, return_log_prob=True,
)
target_q_values = torch.min(
self.target_qf1(next_obs, new_next_actions),
self.target_qf2(next_obs, new_next_actions),
) - alpha * new_log_pi
q_target = self.reward_scale * rewards + (1. - terminals) * self.discount * target_q_values
qf1_loss = self.qf_criterion(q1_pred, q_target.detach())
qf2_loss = self.qf_criterion(q2_pred, q_target.detach())
"""
Policy Loss
"""
qf1_new_actions = self.qf1(obs, new_obs_actions)
qf2_new_actions = self.qf2(obs, new_obs_actions)
q_new_actions = torch.min(
qf1_new_actions,
qf2_new_actions,
)
# Advantage-weighted regression
if self.awr_use_mle_for_vf:
v_pi = self.qf1(obs, policy_mean)
else:
v_pi = self.qf1(obs, new_obs_actions)
if self.awr_sample_actions:
u = new_obs_actions
if self.awr_min_q:
q_adv = q_new_actions
else:
q_adv = qf1_new_actions
else:
u = actions
if self.awr_min_q:
q_adv = torch.min(q1_pred, q2_pred)
else:
q_adv = q1_pred
if self.awr_loss_type == "mse":
policy_logpp = -(policy_mean - actions) ** 2
else:
policy_logpp = dist.log_prob(u)
policy_logpp = policy_logpp.sum(dim=1, keepdim=True)
advantage = q_adv - v_pi
if self.weight_loss and weights is None:
if self.use_automatic_beta_tuning:
_, _, _, _, _, _, _, _, buffer_dist = self.buffer_policy(
obs, reparameterize=True, return_log_prob=True,
)
beta = self.log_beta.exp()
kldiv = torch.distributions.kl.kl_divergence(dist, buffer_dist)
beta_loss = -1*(beta*(kldiv-self.beta_epsilon).detach()).mean()
self.beta_optimizer.zero_grad()
beta_loss.backward()
self.beta_optimizer.step()
else:
beta = self.beta_schedule.get_value(self._n_train_steps_total)
weights = F.softmax(advantage / beta, dim=0)
policy_loss = alpha * log_pi.mean()
if self.use_awr_update and self.weight_loss:
policy_loss = policy_loss + self.awr_weight * (-policy_logpp * len(weights)*weights.detach()).mean()
elif self.use_awr_update:
policy_loss = policy_loss + self.awr_weight * (-policy_logpp).mean()
if self.use_reparam_update:
policy_loss = policy_loss + self.reparam_weight * (-q_new_actions).mean()
policy_loss = self.rl_weight * policy_loss
if self.compute_bc:
train_policy_loss, train_logp_loss, train_mse_loss, _ = self.run_bc_batch(self.demo_train_buffer, self.policy)
policy_loss = policy_loss + self.bc_weight * train_policy_loss
if self.train_bc_on_rl_buffer:
buffer_policy_loss, buffer_train_logp_loss, buffer_train_mse_loss, _ = self.run_bc_batch(self.replay_buffer, self.buffer_policy)
"""
Update networks
"""
if self._n_train_steps_total % self.q_update_period == 0:
self.qf1_optimizer.zero_grad()
qf1_loss.backward()
self.qf1_optimizer.step()
self.qf2_optimizer.zero_grad()
qf2_loss.backward()
self.qf2_optimizer.step()
if self._n_train_steps_total % self.policy_update_period == 0 and self.update_policy:
self.policy_optimizer.zero_grad()
policy_loss.backward()
self.policy_optimizer.step()
if self.train_bc_on_rl_buffer and self._n_train_steps_total % self.policy_update_period == 0 :
self.buffer_policy_optimizer.zero_grad()
buffer_policy_loss.backward()
self.buffer_policy_optimizer.step()
"""
Soft Updates
"""
if self._n_train_steps_total % self.target_update_period == 0:
ptu.soft_update_from_to(
self.qf1, self.target_qf1, self.soft_target_tau
)
ptu.soft_update_from_to(
self.qf2, self.target_qf2, self.soft_target_tau
)
"""
Save some statistics for eval
"""
if self._need_to_update_eval_statistics:
self._need_to_update_eval_statistics = False
"""
Eval should set this to None.
This way, these statistics are only computed for one batch.
"""
policy_loss = (log_pi - q_new_actions).mean()
self.eval_statistics['QF1 Loss'] = np.mean(ptu.get_numpy(qf1_loss))
self.eval_statistics['QF2 Loss'] = np.mean(ptu.get_numpy(qf2_loss))
self.eval_statistics['Policy Loss'] = np.mean(ptu.get_numpy(
policy_loss
))
self.eval_statistics.update(create_stats_ordered_dict(
'Q1 Predictions',
ptu.get_numpy(q1_pred),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Q2 Predictions',
ptu.get_numpy(q2_pred),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Q Targets',
ptu.get_numpy(q_target),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Log Pis',
ptu.get_numpy(log_pi),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Policy mu',
ptu.get_numpy(policy_mean),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Policy log std',
ptu.get_numpy(policy_log_std),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Advantage Weights',
ptu.get_numpy(weights),
))
if self.use_automatic_entropy_tuning:
self.eval_statistics['Alpha'] = alpha.item()
self.eval_statistics['Alpha Loss'] = alpha_loss.item()
if self.compute_bc:
test_policy_loss, test_logp_loss, test_mse_loss, _ = self.run_bc_batch(self.demo_test_buffer, self.policy)
self.eval_statistics.update({
"bc/Train Logprob Loss": ptu.get_numpy(train_logp_loss),
"bc/Test Logprob Loss": ptu.get_numpy(test_logp_loss),
"bc/Train MSE": ptu.get_numpy(train_mse_loss),
"bc/Test MSE": ptu.get_numpy(test_mse_loss),
"bc/train_policy_loss": ptu.get_numpy(train_policy_loss),
"bc/test_policy_loss": ptu.get_numpy(test_policy_loss),
})
if self.train_bc_on_rl_buffer:
test_policy_loss, test_logp_loss, test_mse_loss, _ = self.run_bc_batch(self.replay_buffer,
self.buffer_policy)
_, _, _, _, _, _, _, _, buffer_dist = self.buffer_policy(
obs, reparameterize=True, return_log_prob=True,
)
kldiv = torch.distributions.kl.kl_divergence(dist, buffer_dist)
self.eval_statistics.update({
"buffer_policy/Train Logprob Loss": ptu.get_numpy(buffer_train_logp_loss),
"buffer_policy/Test Logprob Loss": ptu.get_numpy(test_logp_loss),
"buffer_policy/Train MSE": ptu.get_numpy(buffer_train_mse_loss),
"buffer_policy/Test MSE": ptu.get_numpy(test_mse_loss),
"buffer_policy/train_policy_loss": ptu.get_numpy(buffer_policy_loss),
"buffer_policy/test_policy_loss": ptu.get_numpy(test_policy_loss),
"buffer_policy/kl_div":ptu.get_numpy(kldiv.mean()),
})
if self.use_automatic_beta_tuning:
self.eval_statistics.update({
"adaptive_beta/beta":ptu.get_numpy(beta.mean()),
"adaptive_beta/beta loss": ptu.get_numpy(beta_loss.mean()),
})
self._n_train_steps_total += 1
def train_from_torch(self, batch):
logger.push_tabular_prefix("train_q/")
self.eval_statistics = dict()
self._need_to_update_eval_statistics = True
rewards = batch['rewards']
terminals = batch['terminals']
obs = batch['observations']
actions = batch['actions']
next_obs = batch['next_observations']
next_actions = self.target_policy(next_obs)
noise = ptu.randn(next_actions.shape) * self.target_policy_noise
noise = torch.clamp(noise, -self.target_policy_noise_clip,
self.target_policy_noise_clip)
noisy_next_actions = next_actions + noise
target_q1_values = self.target_qf1(next_obs, noisy_next_actions)
target_q2_values = self.target_qf2(next_obs, noisy_next_actions)
target_q_values = torch.min(target_q1_values, target_q2_values)
q_target = self.reward_scale * rewards + (
1. - terminals) * self.discount * target_q_values
q_target = q_target.detach()
q1_pred = self.qf1(obs, actions)
bellman_errors_1 = (q1_pred - q_target)**2
qf1_loss = bellman_errors_1.mean()
q2_pred = self.qf2(obs, actions)
bellman_errors_2 = (q2_pred - q_target)**2
qf2_loss = bellman_errors_2.mean()
"""
Update Networks
"""
self.qf1_optimizer.zero_grad()
qf1_loss.backward()
self.qf1_optimizer.step()
self.qf2_optimizer.zero_grad()
qf2_loss.backward()
self.qf2_optimizer.step()
policy_actions = policy_loss = None
if self._n_train_steps_total % self.policy_update_period == 0:
policy_actions = self.policy(obs)
q_output = self.qf1(obs, policy_actions)
if self.demo_train_buffer._size >= self.bc_batch_size:
if self.use_demo_awr:
train_batch = self.get_batch_from_buffer(
self.demo_train_buffer)
train_o = train_batch["observations"]
train_u = train_batch["actions"]
if self.goal_conditioned:
train_g = train_batch["resampled_goals"]
train_o = torch.cat((train_o, train_g), dim=1)
train_pred_u = self.policy(train_o)
train_error = (train_pred_u - train_u)**2
train_bc_loss = train_error.mean()
policy_q_output_demo_state = self.qf1(
train_o, train_pred_u)
demo_q_output = self.qf1(train_o, train_u)
advantage = demo_q_output - policy_q_output_demo_state
self.eval_statistics['Train BC Loss'] = np.mean(
ptu.get_numpy(train_bc_loss))
if self.awr_policy_update:
train_bc_loss = (train_error * torch.exp(
(advantage) * self.demo_beta))
self.eval_statistics['Advantage'] = np.mean(
ptu.get_numpy(advantage))
if self._n_train_steps_total < self.max_steps_till_train_rl:
rl_weight = 0
else:
rl_weight = self.rl_weight
policy_loss = -rl_weight * q_output.mean(
) + self.bc_weight * train_bc_loss.mean()
else:
train_batch = self.get_batch_from_buffer(
self.demo_train_buffer)
train_o = train_batch["observations"]
# train_pred_u = self.policy(train_o)
if self.goal_conditioned:
train_g = train_batch["resampled_goals"]
train_o = torch.cat((train_o, train_g), dim=1)
train_pred_u = self.policy(train_o)
train_u = train_batch["actions"]
train_error = (train_pred_u - train_u)**2
train_bc_loss = train_error.mean()
# Advantage-weighted regression
policy_error = (policy_actions - actions)**2
policy_error = policy_error.mean(dim=1)
advantage = q1_pred - q_output
weights = F.softmax((advantage / self.beta)[:, 0])
if self.awr_policy_update:
policy_loss = self.rl_weight * (
policy_error * weights.detach() *
self.bc_batch_size).mean()
else:
policy_loss = -self.rl_weight * q_output.mean(
) + self.bc_weight * train_bc_loss.mean()
self.eval_statistics.update(
create_stats_ordered_dict(
'Advantage Weights',
ptu.get_numpy(weights),
))
self.eval_statistics['BC Loss'] = np.mean(
ptu.get_numpy(train_bc_loss))
else: # Normal TD3 update
policy_loss = -self.rl_weight * q_output.mean()
if self.update_policy and not (self.rl_weight == 0
and self.bc_weight == 0):
self.policy_optimizer.zero_grad()
policy_loss.backward()
self.policy_optimizer.step()
self.eval_statistics['Policy Loss'] = np.mean(
ptu.get_numpy(policy_loss))
if self._n_train_steps_total % self.target_update_period == 0:
ptu.soft_update_from_to(self.policy, self.target_policy, self.tau)
ptu.soft_update_from_to(self.qf1, self.target_qf1, self.tau)
ptu.soft_update_from_to(self.qf2, self.target_qf2, self.tau)
if self._need_to_update_eval_statistics:
self._need_to_update_eval_statistics = False
if policy_loss is None:
policy_actions = self.policy(obs)
q_output = self.qf1(obs, policy_actions)
policy_loss = -q_output.mean()
self.eval_statistics['QF1 Loss'] = np.mean(ptu.get_numpy(qf1_loss))
self.eval_statistics['QF2 Loss'] = np.mean(ptu.get_numpy(qf2_loss))
self.eval_statistics.update(
create_stats_ordered_dict(
'Q1 Predictions',
ptu.get_numpy(q1_pred),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Q2 Predictions',
ptu.get_numpy(q2_pred),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Q Targets',
ptu.get_numpy(q_target),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Bellman Errors 1',
ptu.get_numpy(bellman_errors_1),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Bellman Errors 2',
ptu.get_numpy(bellman_errors_2),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Policy Action',
ptu.get_numpy(policy_actions),
))
if self.demo_test_buffer._size >= self.bc_batch_size:
train_batch = self.get_batch_from_buffer(
self.demo_train_buffer)
train_u = train_batch["actions"]
train_o = train_batch["observations"]
if self.goal_conditioned:
train_g = train_batch["resampled_goals"]
train_o = torch.cat((train_o, train_g), dim=1)
train_pred_u = self.policy(train_o)
train_error = (train_pred_u - train_u)**2
train_bc_loss = train_error
policy_q_output_demo_state = self.qf1(train_o, train_pred_u)
demo_q_output = self.qf1(train_o, train_u)
train_advantage = demo_q_output - policy_q_output_demo_state
test_batch = self.get_batch_from_buffer(self.demo_test_buffer)
test_o = test_batch["observations"]
test_u = test_batch["actions"]
if self.goal_conditioned:
test_g = test_batch["resampled_goals"]
test_o = torch.cat((test_o, test_g), dim=1)
test_pred_u = self.policy(test_o)
test_error = (test_pred_u - test_u)**2
test_bc_loss = test_error
policy_q_output_demo_state = self.qf1(test_o, test_pred_u)
demo_q_output = self.qf1(test_o, test_u)
test_advantage = demo_q_output - policy_q_output_demo_state
self.eval_statistics.update(
create_stats_ordered_dict(
'Train BC Loss',
ptu.get_numpy(train_bc_loss),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Train Demo Advantage',
ptu.get_numpy(train_advantage),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Test BC Loss',
ptu.get_numpy(test_bc_loss),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Test Demo Advantage',
ptu.get_numpy(test_advantage),
))
self._n_train_steps_total += 1
logger.pop_tabular_prefix()
def train_from_torch(self, batch):
self._current_epoch += 1
rewards = batch['rewards']
if self._positive_reward:
rewards += 2.5 # Make rewards positive
terminals = batch['terminals']
obs = batch['observations']
actions = batch['actions']
next_obs = batch['next_observations']
"""
Behavior clone a policy
"""
recon, mean, std = self.vae(obs, actions)
recon_loss = self.qf_criterion(recon, actions)
kl_loss = -0.5 * (1 + torch.log(std.pow(2)) - mean.pow(2) -
std.pow(2)).mean()
vae_loss = recon_loss + 0.5 * kl_loss
self.vae_optimizer.zero_grad()
vae_loss.backward()
self.vae_optimizer.step()
"""
Critic Training
"""
# import ipdb; ipdb.set_trace()
with torch.no_grad():
# Duplicate state 10 times (10 is a hyperparameter chosen by BCQ)
state_rep = next_obs.unsqueeze(1).repeat(1, 10, 1).view(
next_obs.shape[0] * 10, next_obs.shape[1])
# Compute value of perturbed actions sampled from the VAE
action_rep = self.policy(state_rep)[0]
target_qf1 = self.target_qf1(state_rep, action_rep)
target_qf2 = self.target_qf2(state_rep, action_rep)
# Soft Clipped Double Q-learning
target_Q = 0.75 * torch.min(target_qf1,
target_qf2) + 0.25 * torch.max(
target_qf1, target_qf2)
target_Q = target_Q.view(next_obs.shape[0],
-1).max(1)[0].view(-1, 1)
target_Q = self.reward_scale * rewards + (
1.0 - terminals) * self.discount * target_Q
qf1_pred = self.qf1(obs, actions)
qf2_pred = self.qf2(obs, actions)
qf1_loss = (qf1_pred - target_Q.detach()).pow(2).mean()
qf2_loss = (qf2_pred - target_Q.detach()).pow(2).mean()
"""
Actor Training
"""
sampled_actions, raw_sampled_actions = self.vae.decode_multiple(
obs, num_decode=self.num_samples_mmd_match)
actor_samples, _, _, _, _, _, _, raw_actor_actions, _ = self.policy(
obs.unsqueeze(1).repeat(1, self.num_samples_mmd_match,
1).view(-1, obs.shape[1]),
return_log_prob=True)
actor_samples = actor_samples.view(obs.shape[0],
self.num_samples_mmd_match,
actions.shape[1])
raw_actor_actions = raw_actor_actions.view(obs.shape[0],
self.num_samples_mmd_match,
actions.shape[1])
if self._use_adv_weighting:
# import ipdb; ipdb.set_trace()
qf1_orig = self.qf1(
obs.unsqueeze(1).repeat(1, self.num_samples_mmd_match,
1).view(-1, obs.shape[1]),
sampled_actions.view(-1, actions.shape[1]))
# Get the target_q for next state
state_rep = next_obs.unsqueeze(1).repeat(1, 10, 1).view(
next_obs.shape[0] * 10, next_obs.shape[1])
action_rep = self.policy(state_rep)[0]
target_q = torch.min(
self.qf1(state_rep, action_rep),
self.qf1(state_rep, action_rep),
)
target_q = target_q.view(next_obs.shape[0],
-1).max(1)[0].view(-1, 1)
target_q = self.reward_scale * rewards + (
1.0 - terminals) * self.discount * target_q
adv_mmd = (
qf1_orig.view(obs.shape[0], self.num_samples_mmd_match, 1) -
target_q.unsqueeze(1))
adv_mmd_not_clipped = adv_mmd
adv_mmd = adv_mmd.exp().clamp_(max=10.0, min=1.0)
adv_mmd = adv_mmd.detach()
if self.kernel_choice == 'laplacian':
if self._use_adv_weighting:
mmd_loss = self.adv_mmd_loss_laplacian(raw_sampled_actions,
raw_actor_actions,
adv_mmd,
sigma=self.mmd_sigma)
else:
mmd_loss = self.mmd_loss_laplacian(raw_sampled_actions,
raw_actor_actions,
sigma=self.mmd_sigma)
elif self.kernel_choice == 'gaussian':
if self._use_adv_weighting:
mmd_loss = self.adv_mmd_loss_gaussian(raw_sampled_actions,
raw_actor_actions,
adv_mmd,
sigma=self.mmd_sigma)
else:
mmd_loss = self.mmd_loss_gaussian(raw_sampled_actions,
raw_actor_actions,
sigma=self.mmd_sigma)
action_divergence = ((sampled_actions - actor_samples)**2).sum(-1)
raw_action_divergence = ((raw_sampled_actions -
raw_actor_actions)**2).sum(-1)
q_val1 = self.qf1(obs, actor_samples[:, 0, :])
q_val2 = self.qf2(obs, actor_samples[:, 0, :])
if self.policy_update_style == '0':
policy_loss = torch.min(q_val1, q_val2)[:, 0]
elif self.policy_update_style == '1':
policy_loss = 0.5 * (q_val1 + q_val2)[:, 0]
if self._n_train_steps_total >= self._bc_pretrain_steps:
# Now we can update the policy
if self.mode == 'auto':
policy_loss = (-policy_loss + self.log_alpha.exp() *
(mmd_loss - self.target_mmd_thresh)).mean()
else:
policy_loss = (-policy_loss + 100 * mmd_loss).mean()
else:
if self.mode == 'auto':
policy_loss = (self.log_alpha.exp() *
(mmd_loss - self.target_mmd_thresh)).mean()
else:
policy_loss = 100 * mmd_loss.mean()
"""
Update Networks
"""
self.qf1_optimizer.zero_grad()
qf1_loss.backward()
self.qf1_optimizer.step()
self.qf2_optimizer.zero_grad()
qf2_loss.backward()
self.qf2_optimizer.step()
self.policy_optimizer.zero_grad()
if self.mode == 'auto':
policy_loss.backward(retain_graph=True)
self.policy_optimizer.step()
if self.mode == 'auto':
self.alpha_optimizer.zero_grad()
(-policy_loss).backward()
self.alpha_optimizer.step()
self.log_alpha.data.clamp_(min=-5.0, max=10.0)
"""
Update networks
"""
if self._use_target_nets:
if self._target_update_method == 'default':
if self._n_train_steps_total % self.target_update_period == 0:
ptu.soft_update_from_to(self.qf1, self.target_qf1,
self.soft_target_tau)
ptu.soft_update_from_to(self.qf2, self.target_qf2,
self.soft_target_tau)
"""
Some statistics for eval
"""
if self._need_to_update_eval_statistics:
self._need_to_update_eval_statistics = False
"""
Eval should set this to None.
This way, these statistics are only computed for one batch.
"""
self.eval_statistics['QF1 Loss'] = np.mean(ptu.get_numpy(qf1_loss))
self.eval_statistics['QF2 Loss'] = np.mean(ptu.get_numpy(qf2_loss))
self.eval_statistics['Num Q Updates'] = self._num_q_update_steps
self.eval_statistics[
'Num Policy Updates'] = self._num_policy_update_steps
if (self._with_gradient_penalty_v1
or self._with_gradient_penalty_v2) and (
self._current_epoch >
self._start_epoch_grad_penalty + 1):
self.eval_statistics['Grad QF1 Loss'] = np.mean(
ptu.get_numpy(grad_qf1_square) * self._grad_coefficient_q)
self.eval_statistics['Grad QF2 Loss'] = np.mean(
ptu.get_numpy(grad_qf2_square) * self._grad_coefficient_q)
self.eval_statistics['Grad Policy Loss'] = np.mean(
ptu.get_numpy(grad_square) * self._grad_coefficient_policy)
self.eval_statistics['Policy Loss'] = np.mean(
ptu.get_numpy(policy_loss))
self.eval_statistics.update(
create_stats_ordered_dict(
'Q1 Predictions',
ptu.get_numpy(qf1_pred),
))
self.eval_statistics.update(
create_stats_ordered_dict(
'Q2 Predictions',
ptu.get_numpy(qf2_pred),
))
if self._use_adv_weighting:
self.eval_statistics.update(
create_stats_ordered_dict(
'Adv MMD', ptu.get_numpy(adv_mmd_not_clipped)))
self.eval_statistics.update(
create_stats_ordered_dict(
'Q Targets',
ptu.get_numpy(target_Q),
))
self.eval_statistics.update(
create_stats_ordered_dict('MMD Loss', ptu.get_numpy(mmd_loss)))
self.eval_statistics.update(
create_stats_ordered_dict('Action Divergence',
ptu.get_numpy(action_divergence)))
self.eval_statistics.update(
create_stats_ordered_dict(
'Raw Action Divergence',
ptu.get_numpy(raw_action_divergence)))
if self.mode == 'auto':
self.eval_statistics['Alpha'] = self.log_alpha.exp().item()
self._n_train_steps_total += 1