def train(self, epoch: int = -1, writer=None) -> str:
self.put_model_on_device()
batch = self.data_loader.get_dataset()
assert len(batch) != 0
values = self._net.critic(inputs=batch.observations, train=False).squeeze().detach()
phi_weights = self._calculate_phi(batch, values).to(self._device).squeeze(-1).detach()
critic_targets = get_reward_to_go(batch).to(self._device) if self._config.phi_key != 'gae' else \
(values + phi_weights).detach()
critic_loss_distribution = self._train_critic_clipped(batch, critic_targets, values)
actor_loss_distribution = self._train_actor_ppo(batch, phi_weights, writer)
if writer is not None:
writer.write_distribution(critic_loss_distribution, "critic_loss")
writer.write_distribution(Distribution(phi_weights.detach()), "phi_weights")
writer.write_distribution(Distribution(critic_targets.detach()), "critic_targets")
if self._config.scheduler_config is not None:
self._actor_scheduler.step()
self._critic_scheduler.step()
self._net.global_step += 1
self.put_model_back_to_original_device()
return f" training policy loss {actor_loss_distribution.mean: 0.3e} [{actor_loss_distribution.std: 0.2e}], " \
f"critic loss {critic_loss_distribution.mean: 0.3e} [{critic_loss_distribution.std: 0.3e}]"
def _train_adversarial_actor_ppo(self, batch: Dataset, phi_weights: torch.Tensor, writer: TensorboardWrapper = None) \
-> Distribution:
original_log_probabilities = self._net.adversarial_policy_log_probabilities(inputs=batch.observations,
actions=batch.actions,
train=False).detach()
list_batch_loss = []
list_entropy_loss = []
for _ in range(self._config.max_actor_training_iterations
if self._config.max_actor_training_iterations != -1 else 1000):
for data in self.data_loader.split_data(np.zeros((0,)), # provide empty array if all data can be selected
batch.observations,
batch.actions,
original_log_probabilities,
phi_weights):
mini_batch_observations, mini_batch_actions, \
mini_batch_original_log_probabilities, mini_batch_phi_weights = data
# normalize advantages (phi_weights)
mini_batch_phi_weights = (mini_batch_phi_weights - mini_batch_phi_weights.mean()) \
/ (mini_batch_phi_weights.std() + 1e-6)
new_log_probabilities = self._net.adversarial_policy_log_probabilities(inputs=mini_batch_observations,
actions=mini_batch_actions,
train=True)
ratio = torch.exp(new_log_probabilities - mini_batch_original_log_probabilities)
unclipped_loss = ratio * mini_batch_phi_weights
clipped_loss = ratio.clamp(1 - self._config.epsilon, 1 + self._config.epsilon) \
* mini_batch_phi_weights
surrogate_loss = - torch.min(unclipped_loss, clipped_loss).mean()
entropy_loss = - self._config.entropy_coefficient * \
self._net.get_adversarial_policy_entropy(mini_batch_observations, train=True).mean()
batch_loss = surrogate_loss + entropy_loss
kl_approximation = (mini_batch_original_log_probabilities - new_log_probabilities).abs().mean().item()
if kl_approximation > 1.5 * self._config.kl_target and self._config.use_kl_stop:
break
self._adversarial_actor_optimizer.zero_grad()
batch_loss.backward()
if self._config.gradient_clip_norm != -1:
nn.utils.clip_grad_norm_(self._net.get_adversarial_actor_parameters(),
self._config.gradient_clip_norm)
self._adversarial_actor_optimizer.step()
assert not np.isnan(batch_loss.detach().numpy())
list_batch_loss.append(batch_loss.detach())
list_entropy_loss.append(entropy_loss.detach())
actor_loss_distribution = Distribution(torch.stack(list_batch_loss))
if writer is not None:
writer.set_step(self._net.global_step)
writer.write_distribution(actor_loss_distribution, "adversarial_policy_loss")
writer.write_distribution(Distribution(torch.stack(list_entropy_loss)), "adversarial_policy_entropy_loss")
writer.write_scalar(list_batch_loss[-1].item(), 'adversarial_final_policy_loss')
writer.write_scalar(kl_approximation, 'adversarial_kl_difference')
return actor_loss_distribution
def train(self, epoch: int = -1, writer=None) -> str:
self.put_model_on_device()
total_error = []
task_error = []
domain_error = []
for source_batch, target_batch in zip(self.data_loader.sample_shuffled_batch(),
self.target_data_loader.sample_shuffled_batch()):
self._optimizer.zero_grad()
targets = data_to_tensor(source_batch.actions).type(self._net.dtype).to(self._device)
# task loss
predictions = self._net.forward(source_batch.observations, train=True)
task_loss = (1 - self._config.epsilon) * self._criterion(predictions, targets).mean()
# add domain adaptation loss
domain_loss = self._config.epsilon * self._domain_adaptation_criterion(
self._net.get_features(source_batch.observations, train=True),
self._net.get_features(target_batch.observations, train=True))
loss = task_loss + domain_loss
loss.backward()
if self._config.gradient_clip_norm != -1:
nn.utils.clip_grad_norm_(self._net.parameters(),
self._config.gradient_clip_norm)
self._optimizer.step()
self._net.global_step += 1
task_error.append(task_loss.cpu().detach())
domain_error.append(domain_loss.cpu().detach())
total_error.append(loss.cpu().detach())
self.put_model_back_to_original_device()
if self._scheduler is not None:
self._scheduler.step()
task_error_distribution = Distribution(task_error)
domain_error_distribution = Distribution(domain_error)
total_error_distribution = Distribution(total_error)
if writer is not None:
writer.set_step(self._net.global_step)
writer.write_distribution(task_error_distribution, 'training/task_error')
writer.write_distribution(domain_error_distribution, 'training/domain_error')
writer.write_distribution(total_error_distribution, 'training/total_error')
if self._config.store_output_on_tensorboard and epoch % 30 == 0:
writer.write_output_image(predictions, 'source/predictions')
writer.write_output_image(targets, 'source/targets')
writer.write_output_image(torch.stack(source_batch.observations), 'source/inputs')
writer.write_output_image(self._net.forward(target_batch.observations, train=True),
'target/predictions')
writer.write_output_image(torch.stack(target_batch.observations), 'target/inputs')
return f' training task: {self._config.criterion} {task_error_distribution.mean: 0.3e} ' \
f'[{task_error_distribution.std:0.2e}]' \
f' domain: {self._config.domain_adaptation_criterion} {domain_error_distribution.mean: 0.3e} ' \
f'[{domain_error_distribution.std:0.2e}]'
def _get_result_message(self, test: bool = False, tag: str = '', lowest_return: bool = False):
msg = f" {'' if not test else 'test'} {tag} "
msg += f"{self._count_episodes} episodes"
if self._count_success != 0:
msg += f" with {self._count_success} success"
if self._writer is not None:
self._writer.write_scalar(self._count_success / float(self._count_episodes), "success")
return_distribution = Distribution(self._episode_returns)
msg += f" with return {return_distribution.mean: 0.3e} [{return_distribution.std: 0.2e}]"
if self._writer is not None:
self._writer.write_scalar(np.mean(self._episode_lengths).item(),
f'{"" if not test else "test_"}episode_lengths{"_"+tag if tag is not "" else ""}')
self._writer.write_distribution(return_distribution,
f'{"" if not test else "test_"}episode_return{"_"+tag if tag is not "" else ""}')
self._writer.write_gif(self._frames,
f'{"" if not test else "test_"}episode{"_"+tag if tag is not "" else ""}')
best_checkpoint = False
if self._max_mean_return is None or return_distribution.mean > self._max_mean_return:
self._max_mean_return = return_distribution.mean
best_checkpoint = True
if self._min_mean_return is None or return_distribution.mean < self._min_mean_return:
self._min_mean_return = return_distribution.mean
if lowest_return:
best_checkpoint = True
return msg, best_checkpoint
def _train_adversarial_critic_clipped(self, batch: Dataset, targets: torch.Tensor, previous_values: torch.Tensor) \
-> Distribution:
critic_loss = []
for value_train_it in range(self._config.max_critic_training_iterations):
state_indices = np.asarray([index for index in range(len(batch)) if not batch.done[index]])
for data in self.data_loader.split_data(state_indices,
batch.observations,
targets,
previous_values):
self._adversarial_critic_optimizer.zero_grad()
mini_batch_observations, mini_batch_targets, mini_batch_previous_values = data
batch_values = self._net.adversarial_critic(inputs=mini_batch_observations, train=True).squeeze()
unclipped_loss = self._criterion(batch_values, mini_batch_targets)
# absolute clipping
clipped_values = mini_batch_previous_values + \
(batch_values - mini_batch_previous_values).clamp(-self._config.epsilon,
self._config.epsilon)
clipped_loss = self._criterion(clipped_values, mini_batch_targets)
batch_loss = torch.max(unclipped_loss, clipped_loss)
batch_loss.mean().backward()
if self._config.gradient_clip_norm != -1:
nn.utils.clip_grad_norm_(self._net.get_adversarial_critic_parameters(),
self._config.gradient_clip_norm)
self._adversarial_critic_optimizer.step()
critic_loss.append(batch_loss.mean().detach())
return Distribution(torch.stack(critic_loss))
def test_distribution_list(self):
data = [10] * 5
distribution = Distribution(data)
self.assertEqual(distribution.max, 10)
self.assertEqual(distribution.min, 10)
self.assertEqual(distribution.mean, 10)
self.assertEqual(distribution.std, 0)
def test_distribution_tensor(self):
data = [10] * 5
distribution = Distribution(torch.as_tensor(data))
self.assertEqual(distribution.max, 10)
self.assertEqual(distribution.min, 10)
self.assertEqual(distribution.mean, 10)
self.assertEqual(distribution.std, 0)
def evaluate(self,
epoch: int = -1,
writer=None,
tag: str = 'validation') -> Tuple[str, bool]:
self.put_model_on_device()
total_error = []
# for batch in tqdm(self.data_loader.get_data_batch(), ascii=True, desc='evaluate'):
for batch in self.data_loader.get_data_batch():
with torch.no_grad():
predictions = self._net.forward(batch.observations,
train=False)
targets = data_to_tensor(batch.actions).type(
self._net.dtype).to(self._device)
error = self._criterion(predictions, targets).mean()
total_error.append(error)
error_distribution = Distribution(total_error)
self.put_model_back_to_original_device()
if writer is not None:
writer.write_distribution(error_distribution, tag)
if self._config.store_output_on_tensorboard and (epoch % 30 == 0
or tag == 'test'):
writer.write_output_image(predictions, f'{tag}/predictions')
writer.write_output_image(targets, f'{tag}/targets')
writer.write_output_image(torch.stack(batch.observations),
f'{tag}/inputs')
msg = f' {tag} {self._config.criterion} {error_distribution.mean: 0.3e} [{error_distribution.std:0.2e}]'
best_checkpoint = False
if self._lowest_validation_loss is None or error_distribution.mean < self._lowest_validation_loss:
self._lowest_validation_loss = error_distribution.mean
best_checkpoint = True
return msg, best_checkpoint
def _train_discriminator_network(self, writer=None) -> str:
total_error = []
criterion = nn.BCELoss()
for sim_batch, real_batch in zip(
self.data_loader.sample_shuffled_batch(),
self.target_data_loader.sample_shuffled_batch()):
self._discriminator_optimizer.zero_grad()
sim_predictions = torch.cat(
self._net.forward_with_all_outputs(sim_batch.observations,
train=False))
real_predictions = torch.cat(
self._net.forward_with_all_outputs(real_batch.observations,
train=False))
targets = torch.as_tensor([*[0] * len(sim_predictions), *[1] * len(real_predictions)])\
.type(self._net.dtype).to(self._device)
outputs = self._net.discriminate(torch.cat(
[sim_predictions, real_predictions]).unsqueeze(dim=1),
train=True).squeeze(dim=1)
loss = criterion(outputs, targets)
loss.mean().backward()
if self._config.gradient_clip_norm != -1:
nn.utils.clip_grad_norm_(self._net.discriminator_parameters(),
self._config.gradient_clip_norm)
self._discriminator_optimizer.step()
total_error.append(loss.cpu().detach())
error_distribution = Distribution(total_error)
if writer is not None:
writer.write_distribution(error_distribution, 'discriminator_loss')
return f' train discriminator network BCE {error_distribution.mean: 0.3e}'
def train(self, epoch: int = -1, writer=None) -> str:
self.put_model_on_device()
total_error = []
for batch in self.data_loader.sample_shuffled_batch():
self._optimizer.zero_grad()
targets = data_to_tensor(batch.actions).type(self._net.dtype).to(
self._device)
probabilities = self._net.forward_with_all_outputs(
batch.observations, train=True)
loss = self._criterion(probabilities[-1], targets).mean()
for index, prob in enumerate(probabilities[:-1]):
loss += self._criterion(prob, targets).mean()
loss.mean().backward()
if self._config.gradient_clip_norm != -1:
nn.utils.clip_grad_norm_(self._net.parameters(),
self._config.gradient_clip_norm)
self._optimizer.step()
self._net.global_step += 1
total_error.append(loss.cpu().detach())
self.put_model_back_to_original_device()
error_distribution = Distribution(total_error)
if writer is not None:
writer.set_step(self._net.global_step)
writer.write_distribution(error_distribution, 'training')
if self._config.store_output_on_tensorboard and epoch % 30 == 0:
for index, prob in enumerate(probabilities):
writer.write_output_image(prob,
f'training/predictions_{index}')
writer.write_output_image(targets, 'training/targets')
writer.write_output_image(torch.stack(batch.observations),
'training/inputs')
return f' training {self._config.criterion} {error_distribution.mean: 0.3e} [{error_distribution.std:0.2e}]'
def train(self, epoch: int = -1, writer=None) -> str:
self.put_model_on_device()
total_error = []
for batch in self.data_loader.sample_shuffled_batch():
self._optimizer.zero_grad()
targets = data_to_tensor(batch.actions).type(self._net.dtype).to(
self._device)
probabilities = self._net.forward_with_all_outputs(
batch.observations, train=True)
loss = self._criterion(probabilities[-1], targets).mean()
for index, prob in enumerate(probabilities[:-1]):
loss += self._criterion(prob, targets).mean()
loss.mean().backward()
if self._config.gradient_clip_norm != -1:
nn.utils.clip_grad_norm_(self._net.parameters(),
self._config.gradient_clip_norm)
self._optimizer.step()
self._net.global_step += 1
total_error.append(loss.cpu().detach())
self.put_model_back_to_original_device()
error_distribution = Distribution(total_error)
if writer is not None:
writer.set_step(self._net.global_step)
writer.write_distribution(error_distribution, 'training')
if self._config.store_output_on_tensorboard and epoch % 30 == 0:
for index, prob in enumerate(probabilities):
writer.write_output_image(prob,
f'training/predictions_{index}')
writer.write_output_image(targets, 'training/targets')
writer.write_output_image(torch.stack(batch.observations),
'training/inputs')
if self._config.store_feature_maps_on_tensorboard and epoch % 30 == 0:
outputs = self._net.forward_with_intermediate_outputs(
batch.observations, train=False)
for i in range(4): # store first 5 images of batch
for layer in ['x1', 'x2', 'x3', 'x4']:
feature_maps = outputs[layer][i].flatten(start_dim=0,
end_dim=0)
title = f'feature_map/layer_{layer}/{i}'
# title += 'inds_' + '_'.join([str(v.item()) for v in winning_indices.indices])
# title += '_vals_' + '_'.join([f'{v.item():0.2f}' for v in winning_indices.values])
writer.write_output_image(feature_maps, title)
writer.write_figure(tag='gradient',
figure=plot_gradient_flow(
self._net.named_parameters()))
return f' training {self._config.criterion} {error_distribution.mean: 0.3e} [{error_distribution.std:0.2e}]'
def train(self, epoch: int = -1, writer=None) -> str:
self.put_model_on_device()
total_error = []
# for batch in tqdm(self.data_loader.sample_shuffled_batch(), ascii=True, desc='train'):
for batch in self.data_loader.sample_shuffled_batch():
self._optimizer.zero_grad()
targets = data_to_tensor(batch.actions).type(self._net.dtype).to(
self._device)
if self._config.add_KL_divergence_loss:
predictions, mean, std = self._net.forward_with_distribution(
batch.observations, train=True)
else:
predictions = self._net.forward(batch.observations, train=True)
loss = self._criterion(predictions, targets).mean()
if self._config.add_KL_divergence_loss:
# https://arxiv.org/pdf/1312.6114.pdf
KL_loss = -0.5 * torch.sum(1 + std.pow(2).log() - mean.pow(2) -
std.pow(2))
loss += KL_loss
loss.backward()
if self._config.gradient_clip_norm != -1:
nn.utils.clip_grad_norm_(self._net.parameters(),
self._config.gradient_clip_norm)
self._optimizer.step()
self._net.global_step += 1
total_error.append(loss.cpu().detach())
self.put_model_back_to_original_device()
if self._scheduler is not None:
self._scheduler.step()
error_distribution = Distribution(total_error)
if writer is not None:
writer.set_step(self._net.global_step)
writer.write_distribution(error_distribution, 'training')
if self._config.add_KL_divergence_loss:
writer.write_scalar(KL_loss, 'KL_divergence')
if self._config.store_output_on_tensorboard and epoch % 30 == 0:
writer.write_output_image(predictions, 'training/predictions')
writer.write_output_image(targets, 'training/targets')
writer.write_output_image(torch.stack(batch.observations),
'training/inputs')
return f' training {self._config.criterion} {error_distribution.mean: 0.3e} [{error_distribution.std:0.2e}]'
def train(self, epoch: int = -1, writer=None) -> str:
self.put_model_on_device()
batch = self.data_loader.get_dataset()
assert len(batch) != 0
values = self._net.critic(inputs=batch.observations,
train=False).squeeze().detach()
phi_weights = self._calculate_phi(batch, values).to(self._device)
policy_loss = self._train_actor(batch, phi_weights)
critic_loss = Distribution(
self._train_critic(batch,
get_reward_to_go(batch).to(self._device)))
if writer is not None:
writer.set_step(self._net.global_step)
writer.write_scalar(policy_loss.data, "policy_loss")
writer.write_distribution(critic_loss, "critic_loss")
self._net.global_step += 1
self.put_model_back_to_original_device()
if self._config.scheduler_config is not None:
self._actor_scheduler.step()
self._critic_scheduler.step()
return f" training policy loss {policy_loss.data: 0.3e}, critic loss {critic_loss.mean: 0.3e}"
def train(self, epoch: int = -1, writer=None) -> str:
self.put_model_on_device()
total_error = []
task_error = []
domain_error = []
for source_batch, target_batch in zip(self.data_loader.sample_shuffled_batch(),
self.target_data_loader.sample_shuffled_batch()):
self._optimizer.zero_grad()
targets = data_to_tensor(source_batch.actions).type(self._net.dtype).to(self._device)
# deep supervision loss
probabilities = self._net.forward_with_all_outputs(source_batch.observations, train=True)
task_loss = self._criterion(probabilities[-1], targets).mean()
for index, prob in enumerate(probabilities[:-1]):
task_loss += self._criterion(prob, targets).mean()
task_loss *= (1 - self._config.epsilon)
# add domain adaptation loss on distribution of output pixels at each output
domain_loss = sum([self._domain_adaptation_criterion(sp.flatten().unsqueeze(1), tp.flatten().unsqueeze(1))
for sp, tp in zip(self._net.forward_with_all_outputs(source_batch.observations,
train=True),
self._net.forward_with_all_outputs(target_batch.observations,
train=True))
]) * self._config.epsilon
loss = task_loss + domain_loss
loss.backward()
if self._config.gradient_clip_norm != -1:
nn.utils.clip_grad_norm_(self._net.parameters(),
self._config.gradient_clip_norm)
self._optimizer.step()
self._net.global_step += 1
task_error.append(task_loss.cpu().detach().numpy())
domain_error.append(domain_loss.cpu().detach().numpy())
total_error.append(loss.cpu().detach().numpy())
self.put_model_back_to_original_device()
if self._scheduler is not None:
self._scheduler.step()
task_error_distribution = Distribution(task_error)
domain_error_distribution = Distribution(domain_error)
total_error_distribution = Distribution(total_error)
if writer is not None:
writer.set_step(self._net.global_step)
writer.write_distribution(task_error_distribution, 'training/task_error')
writer.write_distribution(domain_error_distribution, 'training/domain_error')
writer.write_distribution(total_error_distribution, 'training/total_error')
if self._config.store_output_on_tensorboard and epoch % 30 == 0:
writer.write_output_image(probabilities[-1], 'source/predictions')
writer.write_output_image(targets, 'source/targets')
writer.write_output_image(torch.stack(source_batch.observations), 'source/inputs')
writer.write_output_image(self._net.forward(target_batch.observations, train=False),
'target/predictions')
writer.write_output_image(torch.stack(target_batch.observations), 'target/inputs')
if self._config.store_feature_maps_on_tensorboard and epoch % 30 == 0:
for name, batch in zip(['source', 'target'], [source_batch, target_batch]):
outputs = self._net.forward_with_intermediate_outputs(batch.observations, train=False)
for i in range(4): # store first 5 images of batch
for layer in ['x1', 'x2', 'x3', 'x4']:
feature_maps = outputs[layer][i].flatten(start_dim=0, end_dim=0)
title = f'feature_map/{name}/layer_{layer}/{i}'
# title += 'inds_' + '_'.join([str(v.item()) for v in winning_indices.indices])
# title += '_vals_' + '_'.join([f'{v.item():0.2f}' for v in winning_indices.values])
writer.write_output_image(feature_maps, title)
return f' task {self._config.criterion} ' \
f'{task_error_distribution.mean: 0.3e} ' \
f'[{task_error_distribution.std:0.2e}] ' \
f' domain {self._config.domain_adaptation_criterion} ' \
f'{domain_error_distribution.mean: 0.3e} ' \
f'[{domain_error_distribution.std: 0.2e}]'
def _train_main_network(self, epoch: int = -1, writer=None) -> str:
deeply_supervised_error = []
discriminator_error = []
for sim_batch, real_batch in zip(
self.data_loader.sample_shuffled_batch(),
self.target_data_loader.sample_shuffled_batch()):
self._optimizer.zero_grad()
# normal deep supervision loss
targets = data_to_tensor(sim_batch.actions).type(
self._net.dtype).to(self._device)
probabilities = self._net.forward_with_all_outputs(
sim_batch.observations, train=True)
loss = self._criterion(probabilities[-1], targets).mean()
for index, prob in enumerate(probabilities[:-1]):
loss += self._criterion(prob, targets).mean()
deeply_supervised_error.append(loss.mean().cpu().detach())
# adversarial loss on discriminator data
network_outputs = torch.cat(
self._net.forward_with_all_outputs(
real_batch.observations, train=True)).unsqueeze(dim=1)
discriminator_loss = self._net.discriminate(network_outputs,
train=False).mean()
#results = self._net.forward_with_intermediate_outputs(real_batch.observations, train=True)
#feature_maps =
# combine losses with epsilon weight
loss *= (1 - self._config.epsilon)
loss += self._config.epsilon * discriminator_loss
loss.mean().backward()
discriminator_error.append(
discriminator_loss.mean().cpu().detach())
# clip gradients
if self._config.gradient_clip_norm != -1:
nn.utils.clip_grad_norm_(self._net.parameters(),
self._config.gradient_clip_norm)
self._optimizer.step()
supervised_error_distribution = Distribution(deeply_supervised_error)
discriminator_error_distribution = Distribution(discriminator_error)
if writer is not None:
writer.write_distribution(supervised_error_distribution,
'training_loss_from_deep_supervision')
writer.write_distribution(discriminator_error_distribution,
'training_loss_from_discriminator')
if self._config.store_output_on_tensorboard and epoch % 30 == 0:
for index, prob in enumerate(probabilities):
writer.write_output_image(prob,
f'training/predictions_{index}')
writer.write_output_image(targets, 'training/targets')
writer.write_output_image(torch.stack(sim_batch.observations),
'training/inputs')
for index, prob in enumerate(
self._net.forward_with_all_outputs(real_batch.observations,
train=False)):
writer.write_output_image(prob, f'real/predictions_{index}')
writer.write_output_image(torch.stack(real_batch.observations),
'real/inputs')
return f' Training: supervision {self._config.criterion} {supervised_error_distribution.mean: 0.3e} ' \
f'[{supervised_error_distribution.std:0.2e}]' \
f' discriminator {discriminator_error_distribution.mean: 0.3e} ' \
f'[{discriminator_error_distribution.std:0.2e}]'