def __init__( self, train_dataset, test_dataset, model, batch_size=128, log_interval=0, beta=0.5, beta_schedule=None, imsize=84, lr=1e-3, do_scatterplot=False, normalize=False, state_sim_debug=False, mse_weight=0.1, is_auto_encoder=False, lmbda=0.5, mu=1, gamma=0.2, ): self.log_interval = log_interval self.batch_size = batch_size self.beta = beta if is_auto_encoder: self.beta = 0 self.beta_schedule = beta_schedule if self.beta_schedule is None: self.beta_schedule = ConstantSchedule(self.beta) self.imsize = imsize self.do_scatterplot = do_scatterplot self.lmbda = lmbda self.mu = mu self.gamma = gamma """ I think it's a bit nicer if the caller makes this call, i.e. ``` m = ConvVAE(representation_size) if ptu.gpu_enabled(): m.cuda() t = ConvVAETrainer(train_data, test_data, m) ``` However, I'll leave this here for backwards-compatibility. """ if ptu.gpu_enabled(): model.cuda() self.model = model self.representation_size = model.representation_size self.input_channels = model.input_channels self.imlength = model.imlength self.optimizer = optim.Adam(self.model.parameters(), lr=lr) self.train_dataset, self.test_dataset = train_dataset, test_dataset self.normalize = normalize self.state_sim_debug = state_sim_debug self.mse_weight = mse_weight self.x_next_index = self.input_channels * self.imsize**2 if self.normalize: self.train_data_mean = np.mean(self.train_dataset, axis=0)
def __init__( self, model, batch_size=128, log_interval=0, beta=0.5, beta_schedule=None, lr=None, weight_decay=0, ): self.model = model self.log_interval = log_interval self.batch_size = batch_size self.beta = beta if lr is None: if is_auto_encoder: lr = 1e-2 else: lr = 1e-3 self.beta_schedule = beta_schedule if self.beta_schedule is None or is_auto_encoder: self.beta_schedule = ConstantSchedule(self.beta) self.imsize = model.imsize model.to(ptu.device) self.representation_size = model.representation_size self.input_channels = model.input_channels self.imlength = self.imsize * self.imsize * self.input_channels self.lr = lr params = list(self.model.parameters()) self.optimizer = optim.Adam( params, lr=self.lr, weight_decay=weight_decay, ) self.eval_statistics = {}
def __init__( self, max_tau=10, epoch_max_tau_schedule=None, sample_train_goals_from='replay_buffer', sample_rollout_goals_from='environment', cycle_taus_for_rollout=False, ): """ :param max_tau: Maximum tau (planning horizon) to train with. :param epoch_max_tau_schedule: A schedule for the maximum planning horizon tau. :param sample_train_goals_from: Sampling strategy for goals used in training. Can be one of the following strings: - environment: Sample from the environment - replay_buffer: Sample from the replay_buffer - her: Sample from a HER-based replay_buffer :param sample_rollout_goals_from: Sampling strategy for goals used during rollout. Can be one of the following strings: - environment: Sample from the environment - replay_buffer: Sample from the replay_buffer - fixed: Do no resample the goal. Just use the one in the environment. :param vectorized: Train the QF in vectorized form? :param cycle_taus_for_rollout: Decrement the tau passed into the policy during rollout? """ assert sample_train_goals_from in ['environment', 'replay_buffer', 'her'] assert sample_rollout_goals_from in ['environment', 'replay_buffer', 'fixed'] if epoch_max_tau_schedule is None: epoch_max_tau_schedule = ConstantSchedule(max_tau) self.max_tau = max_tau self.epoch_max_tau_schedule = epoch_max_tau_schedule self.sample_train_goals_from = sample_train_goals_from self.sample_rollout_goals_from = sample_rollout_goals_from self.cycle_taus_for_rollout = cycle_taus_for_rollout self._current_path_goal = None self._rollout_tau = self.max_tau self.policy = MakeUniversal(self.policy) self.eval_policy = MakeUniversal(self.eval_policy) self.exploration_policy = MakeUniversal(self.exploration_policy) self.eval_sampler = MultigoalSimplePathSampler( env=self.env, policy=self.eval_policy, max_samples=self.num_steps_per_eval, max_path_length=self.max_path_length, discount_sampling_function=self._sample_max_tau_for_rollout, goal_sampling_function=self._sample_goal_for_rollout, cycle_taus_for_rollout=self.cycle_taus_for_rollout, ) if self.collection_mode == 'online-parallel': # TODO(murtaza): What happens to the eval env? # see `eval_sampler` definition above. self.training_env = RemoteRolloutEnv( env=self.env, policy=self.eval_policy, exploration_policy=self.exploration_policy, max_path_length=self.max_path_length, normalize_env=self.normalize_env, rollout_function=self.rollout, )
def train_gan(variant, return_data=False): from rlkit.torch.gan.dcgan import Generator, Discriminator from rlkit.torch.gan.dcgan_trainer import DCGANTrainer from rlkit.torch.gan.bigan import Generator, Encoder, Discriminator from rlkit.torch.gan.bigan_trainer import BiGANTrainer from rlkit.misc.ml_util import PiecewiseLinearSchedule, ConstantSchedule from rlkit.core import logger import rlkit.torch.pytorch_util as ptu from rlkit.pythonplusplus import identity import torch import torch.utils.data import torchvision.datasets as dset from rlkit.data_management.external.bair_dataset import bair_dataset import torchvision.transforms as transforms from rlkit.data_management.external.bair_dataset.config import BAIR_DATASET_LOCATION from rlkit.misc.asset_loader import sync_down_folder if not variant.get('simpusher', False): if variant["dataset"] == "bair": #train_dataset, test_dataset, info dataloader = bair_dataset.generate_dataset( variant['generate_dataset_kwargs'])[0].dataset_loader get_data = lambda d: d['x_t'] if variant["dataset"] == "cifar10": local_path = sync_down_folder(variant["dataroot"]) #local_path = variant["dataroot"] dataset = dset.CIFAR10(root=local_path, train=True, download=False, transform=transforms.Compose( [transforms.ToTensor()])) dataloader = torch.utils.data.DataLoader( dataset, batch_size=variant["batch_size"], shuffle=True, num_workers=variant["num_workers"]) get_data = lambda d: d[0] if variant["dataset"] == "celebfaces": local_path = sync_down_folder(variant["dataroot"]) #local_path = variant["dataroot"] dataset = dset.ImageFolder( root=local_path, transform=transforms.Compose([ transforms.Resize(variant["image_size"]), transforms.CenterCrop(variant["image_size"]), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ])) dataloader = torch.utils.data.DataLoader( dataset, batch_size=variant["batch_size"], shuffle=True, num_workers=variant["num_workers"]) get_data = lambda d: d[0] model_class = variant["gan_class"] trainer_class = variant["gan_trainer_class"] if trainer_class is DCGANTrainer: model = model_class(variant["ngpu"], variant["nc"], variant["latent_size"], variant["ngf"], variant["ndf"]) trainer = trainer_class(model, variant["lr"], variant["beta"], variant["latent_size"]) if trainer_class is BiGANTrainer: model = model_class(variant["ngpu"], variant["latent_size"], variant["dropout"], variant["output_size"]) trainer = trainer_class(model, variant["ngpu"], variant["lr"], variant["beta"], variant["latent_size"], variant["generator_threshold"]) if variant.get('simpusher', False): imsize = variant.get('imsize') beta = variant["beta"] representation_size = variant.get("representation_size", variant.get("latent_size", None)) use_linear_dynamics = variant.get('use_linear_dynamics', False) generate_vae_dataset_fctn = variant.get('generate_vae_data_fctn', generate_vae_dataset) variant['generate_vae_dataset_kwargs'][ 'use_linear_dynamics'] = use_linear_dynamics batch_size = variant['algo_kwargs']['batch_size'] variant['generate_vae_dataset_kwargs']['batch_size'] = batch_size train_dataset, test_dataset, info = generate_vae_dataset_fctn( variant['generate_vae_dataset_kwargs']) dataloader = train_dataset.dataset_loader get_data = lambda d: d['x_t'].reshape(128, 3, imsize, imsize) if use_linear_dynamics: action_dim = train_dataset.data['actions'].shape[2] logger.save_extra_data(info) logger.get_snapshot_dir() if 'context_schedule' in variant: schedule = variant['context_schedule'] if type(schedule) is dict: context_schedule = PiecewiseLinearSchedule(**schedule) else: context_schedule = ConstantSchedule(schedule) variant['algo_kwargs']['context_schedule'] = context_schedule if variant['algo_kwargs'].get('is_auto_encoder', False): model = AutoEncoder(representation_size, **variant['gan_kwargs']) elif variant.get('use_spatial_auto_encoder', False): model = SpatialAutoEncoder(representation_size, **variant['gan_kwargs']) else: gan_class = variant['gan_class'] if use_linear_dynamics: model = gan_class(latent_size=representation_size, **variant['gan_kwargs']) else: model = gan_class(latent_size=representation_size, **variant['gan_kwargs']) model.to(ptu.device) gan_trainer_class = variant['gan_trainer_class'] trainer = gan_trainer_class(model, latent_size=representation_size, beta=beta, **variant['algo_kwargs']) save_period = variant['save_period'] dump_skew_debug_plots = variant.get('dump_skew_debug_plots', False) for epoch in range(variant['num_epochs']): trainer.train_epoch(dataloader, epoch, variant['num_epochs'], get_data) #trainer.test_epoch(epoch, test_dataset) #dump samples is called in trainer stats = trainer.get_stats(epoch) for k, v in stats.items(): logger.record_tabular(k, v) logger.dump_tabular() #trainer.end_epoch(epoch) if epoch % 50 == 0: logger.save_itr_params(epoch, trainer.get_model()) logger.save_extra_data(trainer.get_model(), 'gan.pkl', mode='pickle') if return_data: return model, train_dataset, test_dataset return model
class ACAITrainer(): def __init__( self, train_dataset, test_dataset, model, batch_size=128, log_interval=0, beta=0.5, beta_schedule=None, imsize=84, lr=1e-3, do_scatterplot=False, normalize=False, state_sim_debug=False, mse_weight=0.1, is_auto_encoder=False, lmbda=0.5, mu=1, gamma=0.2, ): self.log_interval = log_interval self.batch_size = batch_size self.beta = beta if is_auto_encoder: self.beta = 0 self.beta_schedule = beta_schedule if self.beta_schedule is None: self.beta_schedule = ConstantSchedule(self.beta) self.imsize = imsize self.do_scatterplot = do_scatterplot self.lmbda = lmbda self.mu = mu self.gamma = gamma """ I think it's a bit nicer if the caller makes this call, i.e. ``` m = ConvVAE(representation_size) if ptu.gpu_enabled(): m.cuda() t = ConvVAETrainer(train_data, test_data, m) ``` However, I'll leave this here for backwards-compatibility. """ if ptu.gpu_enabled(): model.cuda() self.model = model self.representation_size = model.representation_size self.input_channels = model.input_channels self.imlength = model.imlength self.optimizer = optim.Adam(self.model.parameters(), lr=lr) self.train_dataset, self.test_dataset = train_dataset, test_dataset self.normalize = normalize self.state_sim_debug = state_sim_debug self.mse_weight = mse_weight self.x_next_index = self.input_channels * self.imsize**2 if self.normalize: self.train_data_mean = np.mean(self.train_dataset, axis=0) # self.train_dataset = ((self.train_dataset - self.train_data_mean)) + 1 / 2 # self.test_dataset = ((self.test_dataset - self.train_data_mean)) + 1 / 2 def get_batch(self, train=True): dataset = self.train_dataset if train else self.test_dataset ind = np.random.randint(0, len(dataset), self.batch_size) samples = dataset[ind, :] samples = normalize_image(samples) if self.normalize: samples = ((samples - self.train_data_mean) + 1) / 2 return ptu.np_to_var(samples) def get_debug_batch(self, train=True): dataset = self.train_dataset if train else self.test_dataset X, Y = dataset ind = np.random.randint(0, Y.shape[0], self.batch_size) X = X[ind, :] Y = Y[ind, :] return ptu.np_to_var(X), ptu.np_to_var(Y) def get_batch_smooth(self, train=True): dataset = self.train_dataset if train else self.test_dataset ind = np.random.randint(0, len(dataset), self.batch_size) samples = dataset[ind, :] samples = normalize_image(samples) if self.normalize: samples = ((samples - self.train_data_mean) + 1) / 2 x_next, x = samples[:, :self.x_next_index], samples[:, self.x_next_index:] return ptu.np_to_var(x_next), ptu.np_to_var(x) def logprob(self, recon_x, x, mu, logvar): # Divide by batch_size rather than setting size_average=True because # otherwise the averaging will also happen across dimension 1 (the # pixels) return F.binary_cross_entropy( recon_x, x.narrow(start=0, length=self.imlength, dimension=1).contiguous().view(-1, self.imlength), size_average=False, ) / self.batch_size def kl_divergence(self, recon_x, x, mu, logvar): return -torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1).mean() def state_similarity_loss(self, model, encoded_x, states): output = self.model.fc6(F.relu(self.model.fc5(encoded_x))) return torch.norm(output - states)**2 / self.batch_size def train_epoch(self, epoch, sample_batch=None, batches=100, from_rl=False): self.model.train() losses = [] bces = [] kles = [] mses = [] losses_c = [] beta = self.beta_schedule.get_value(epoch) for batch_idx in range(batches): data = self.get_batch() if sample_batch is not None: data = sample_batch(self.batch_size) self.optimizer.zero_grad() recon_batch, mu, logvar, predicted_alpha, alpha = self.model(data) bce = self.logprob(recon_batch, data, mu, logvar) kle = self.kl_divergence(recon_batch, data, mu, logvar) loss = bce + beta * kle + self.lmbda * torch.norm( predicted_alpha, 2) regularizer_a = self.model.critic(self.gamma * data + (1 - self.gamma) * recon_batch) loss_c = torch.norm(predicted_alpha - alpha, 2) + self.mu * torch.norm(regularizer_a, 2) loss.backward(retain_graph=True) loss_c.backward() losses.append(loss.data[0]) losses_c.append(loss_c.data[0]) bces.append(bce.data[0]) kles.append(kle.data[0]) if self.state_sim_debug: mses.append(sim_loss.data[0]) self.optimizer.step() if self.log_interval and batch_idx % self.log_interval == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, batch_idx * len(data), len(self.train_loader.dataset), 100. * batch_idx / len(self.train_loader), loss.data[0] / len(data))) if not from_rl: logger.record_tabular("train/epoch", epoch) logger.record_tabular("train/critic_loss", np.mean(losses_c)) logger.record_tabular("train/BCE", np.mean(bces)) logger.record_tabular("train/KL", np.mean(kles)) if self.state_sim_debug: logger.record_tabular("train/mse", np.mean(mses)) logger.record_tabular("train/loss", np.mean(losses)) def test_epoch( self, epoch, save_reconstruction=True, save_scatterplot=True, save_vae=True, from_rl=False, ): self.model.eval() losses = [] losses_c = [] bces = [] kles = [] zs = [] mses = [] beta = self.beta_schedule.get_value(epoch) for batch_idx in range(10): data = self.get_batch(train=False) recon_batch, mu, logvar, predicted_alpha, alpha = self.model(data) bce = self.logprob(recon_batch, data, mu, logvar) kle = self.kl_divergence(recon_batch, data, mu, logvar) loss = bce + beta * kle + self.lmbda * torch.norm( predicted_alpha, 2) regularizer_a = self.model.critic(self.gamma * data + (1 - self.gamma) * recon_batch) loss_c = torch.norm(predicted_alpha - alpha, 2) + self.mu * torch.norm(regularizer_a, 2) z_data = ptu.get_numpy(mu.cpu()) for i in range(len(z_data)): zs.append(z_data[i, :]) losses.append(loss.data[0]) losses_c.append(loss_c.data[0]) bces.append(bce.data[0]) kles.append(kle.data[0]) if self.state_sim_debug: mses.append(sim_loss.data[0]) if batch_idx == 0 and save_reconstruction: n = min(data.size(0), 8) comparison = torch.cat([ data[:n].narrow(start=0, length=self.imlength, dimension=1).contiguous().view( -1, self.input_channels, self.imsize, self.imsize), recon_batch.view( self.batch_size, self.input_channels, self.imsize, self.imsize, )[:n] ]) save_dir = osp.join(logger.get_snapshot_dir(), 'r%d.png' % epoch) save_image(comparison.data.cpu(), save_dir, nrow=n) zs = np.array(zs) self.model.dist_mu = zs.mean(axis=0) self.model.dist_std = zs.std(axis=0) if self.do_scatterplot and save_scatterplot: self.plot_scattered(np.array(zs), epoch) if not from_rl: for key, value in self.debug_statistics().items(): logger.record_tabular(key, value) logger.record_tabular("test/BCE", np.mean(bces)) logger.record_tabular("test/KL", np.mean(kles)) logger.record_tabular("test/loss", np.mean(losses)) logger.record_tabular("test/critic_loss", np.mean(losses_c)) logger.record_tabular("beta", beta) if self.state_sim_debug: logger.record_tabular("test/MSE", np.mean(mses)) logger.dump_tabular() if save_vae: logger.save_itr_params(epoch, self.model) # slow... # logdir = logger.get_snapshot_dir() # filename = osp.join(logdir, 'params.pkl') # torch.save(self.model, filename) def debug_statistics(self): """ Given an image $$x$$, samples a bunch of latents from the prior $$z_i$$ and decode them $$\hat x_i$$. Compare this to $$\hat x$$, the reconstruction of $$x$$. Ideally - All the $$\hat x_i$$s do worse than $$\hat x$$ (makes sure VAE isn’t ignoring the latent) - Some $$\hat x_i$$ do better than other $$\hat x_i$$ (tests for coverage) """ debug_batch_size = 64 data = self.get_batch(train=False) recon_batch, mu, logvar, predicted_alpha, alpha = self.model(data) img = data[0] recon_mse = ((recon_batch[0] - img)**2).mean() img_repeated = img.expand((debug_batch_size, img.shape[0])) samples = ptu.Variable( torch.randn(debug_batch_size, self.representation_size)) random_imgs = self.model.decode(samples) random_mse = ((random_imgs - img_repeated)**2).mean(dim=1) mse_improvement = ptu.get_numpy(random_mse - recon_mse) stats = create_stats_ordered_dict( 'debug/MSE improvement over random', mse_improvement, ) stats['debug/MSE of random reconstruction'] = ptu.get_numpy( recon_mse)[0] return stats def dump_samples(self, epoch): self.model.eval() sample = ptu.Variable(torch.randn(64, self.representation_size)) sample = self.model.decode(sample).cpu() save_dir = osp.join(logger.get_snapshot_dir(), 's%d.png' % epoch) save_image( sample.data.view(64, self.input_channels, self.imsize, self.imsize), save_dir) def plot_scattered(self, z, epoch): try: import matplotlib.pyplot as plt except ImportError: logger.log(__file__ + ": Unable to load matplotlib. Consider " "setting do_scatterplot to False") return dim_and_stds = [(i, np.std(z[:, i])) for i in range(z.shape[1])] dim_and_stds = sorted(dim_and_stds, key=lambda x: x[1]) dim1 = dim_and_stds[-1][0] dim2 = dim_and_stds[-2][0] plt.figure(figsize=(8, 8)) plt.scatter(z[:, dim1], z[:, dim2], marker='o', edgecolor='none') if self.model.dist_mu is not None: x1 = self.model.dist_mu[dim1:dim1 + 1] y1 = self.model.dist_mu[dim2:dim2 + 1] x2 = self.model.dist_mu[dim1:dim1 + 1] + self.model.dist_std[dim1:dim1 + 1] y2 = self.model.dist_mu[dim2:dim2 + 1] + self.model.dist_std[dim2:dim2 + 1] plt.plot([x1, x2], [y1, y2], color='k', linestyle='-', linewidth=2) axes = plt.gca() axes.set_xlim([-6, 6]) axes.set_ylim([-6, 6]) axes.set_title('dim {} vs dim {}'.format(dim1, dim2)) plt.grid(True) save_file = osp.join(logger.get_snapshot_dir(), 'scatter%d.png' % epoch) plt.savefig(save_file)
def __init__( self, env, policy, qf1, qf2, target_qf1, target_qf2, reward_type, discount_type=None, discount_model=None, discount=0.99, prior_discount_weight_schedule: Optional[ScalarSchedule] = None, multiply_bootstrap_by_prior_discount=False, upper_bound_discount_by_prior=False, reward_scale=1.0, reward_tracking_momentum=0.999, auto_init_qf_bias=False, policy_lr=1e-3, qf_lr=1e-3, optimizer_class=optim.Adam, soft_target_tau=1e-2, target_update_period=1, plotter=None, render_eval_paths=False, use_automatic_entropy_tuning=True, target_entropy=None, ): """ :param env: :param policy: :param qf1: :param qf2: :param target_qf1: :param target_qf2: :param reward_type: :param discount_type: :param discount_model: :param discount: :param prior_discount_weight_schedule: At epoch i, use the discount discount = c_i * prior_discount + (1-c_i) posterior_discount :param multiply_bootstrap_by_prior_discount: If true, when you compute the discount prior, always multiply it by the prior discount in addition to the normal prior discount. :param upper_bound_discount_by_prior: Always upper-bound the discount by the prior. :param reward_scale: :param reward_tracking_momentum: :param policy_lr: :param qf_lr: :param optimizer_class: :param soft_target_tau: :param target_update_period: :param plotter: :param render_eval_paths: :param use_automatic_entropy_tuning: :param target_entropy: """ if reward_type not in { self.NORMAL_REWARD, self.DISCOUNTED_REWARD, self.DISCOUNTED_PLUS_TIME_KL, }: raise ValueError("Invalid reward type: {}".format(reward_type)) if discount_type is None: # preserve old behavior if reward_type == self.DISCOUNTED_PLUS_TIME_KL: discount_type = self.COMPUTED_DISCOUNT else: discount_type = self.PRIOR_DISCOUNT if discount_type not in { self.PRIOR_DISCOUNT, self.LEARNED_DISCOUNT, self.COMPUTED_DISCOUNT, self.COMPUTED_DISCOUNT_NO_PRIOR, }: raise ValueError("Invalid discount type: {}".format(discount_type)) if (reward_type == self.LEARNED_DISCOUNT and discount_model is None): raise ValueError( "Need to set discount_model for using mode {}".format( reward_type)) if not isinstance(reward_scale, Number) and reward_scale not in { 'auto_normalize_by_max_magnitude', 'auto_normalize_by_max_magnitude_times_10', 'auto_normalize_by_max_magnitude_times_100', 'auto_normalize_by_max_magnitude_times_invsig_prior', 'auto_normalize_by_mean_magnitude', }: raise ValueError( "Invalid reward_scale type: {}".format(reward_scale)) super().__init__() self.env = env self.policy = policy self.qf1 = qf1 self.qf2 = qf2 self.target_qf1 = target_qf1 self.target_qf2 = target_qf2 self.soft_target_tau = soft_target_tau self.target_update_period = target_update_period self.reward_type = reward_type self.discount_type = discount_type if prior_discount_weight_schedule is None: prior_discount_weight_schedule = ConstantSchedule(0.) self._prior_discount_weight_schedule = prior_discount_weight_schedule self._multiply_bootstrap_by_prior_discount = ( multiply_bootstrap_by_prior_discount) self._upper_bound_discount_by_prior = (upper_bound_discount_by_prior) self._auto_init_qf_bias = auto_init_qf_bias self._current_epoch = 0 self.use_automatic_entropy_tuning = use_automatic_entropy_tuning if self.use_automatic_entropy_tuning: if target_entropy is None: # Use heuristic value from SAC paper self.target_entropy = -np.prod( self.env.action_space.shape).item() else: self.target_entropy = target_entropy self.log_alpha = ptu.zeros(1, requires_grad=True) self.alpha_optimizer = optimizer_class( [self.log_alpha], lr=policy_lr, ) self.plotter = plotter self.render_eval_paths = render_eval_paths self.qf_criterion = nn.MSELoss() self.vf_criterion = nn.MSELoss() self.policy_optimizer = optimizer_class( self.policy.parameters(), lr=policy_lr, ) self.qf1_optimizer = optimizer_class( self.qf1.parameters(), lr=qf_lr, ) self.qf2_optimizer = optimizer_class( self.qf2.parameters(), lr=qf_lr, ) self.discount_model = discount_model self.discount = discount self.prior_on_discount = Bernoulli(self.discount) self._reward_scale = reward_scale self.reward_tracking_momentum = reward_tracking_momentum self._reward_normalizer = ptu.from_numpy(np.array(1)) self.eval_statistics = OrderedDict() self._need_to_update_eval_statistics = True self._qfs_were_initialized = False
def __init__( self, env, policy, qf1, qf2, target_qf1, target_qf2, buffer_policy=None, discount=0.99, reward_scale=1.0, beta=1.0, beta_schedule_kwargs=None, policy_lr=1e-3, qf_lr=1e-3, policy_weight_decay=0, q_weight_decay=0, optimizer_class=optim.Adam, soft_target_tau=1e-2, target_update_period=1, plotter=None, render_eval_paths=False, use_automatic_entropy_tuning=True, target_entropy=None, bc_num_pretrain_steps=0, q_num_pretrain1_steps=0, q_num_pretrain2_steps=0, bc_batch_size=128, alpha=1.0, policy_update_period=1, q_update_period=1, weight_loss=True, compute_bc=True, use_awr_update=True, use_reparam_update=False, bc_weight=0.0, rl_weight=1.0, reparam_weight=1.0, awr_weight=1.0, post_pretrain_hyperparams=None, post_bc_pretrain_hyperparams=None, awr_use_mle_for_vf=False, vf_K=1, awr_sample_actions=False, buffer_policy_sample_actions=False, awr_min_q=False, brac=False, reward_transform_class=None, reward_transform_kwargs=None, terminal_transform_class=None, terminal_transform_kwargs=None, pretraining_logging_period=1000, train_bc_on_rl_buffer=False, use_automatic_beta_tuning=False, beta_epsilon=1e-10, normalize_over_batch=True, normalize_over_state="advantage", Z_K=10, clip_score=None, validation_qlearning=False, mask_positive_advantage=False, buffer_policy_reset_period=-1, num_buffer_policy_train_steps_on_reset=100, advantage_weighted_buffer_loss=True, ): super().__init__() self.env = env self.policy = policy self.qf1 = qf1 self.qf2 = qf2 self.target_qf1 = target_qf1 self.target_qf2 = target_qf2 self.buffer_policy = buffer_policy self.soft_target_tau = soft_target_tau self.target_update_period = target_update_period self.use_awr_update = use_awr_update self.use_automatic_entropy_tuning = use_automatic_entropy_tuning if self.use_automatic_entropy_tuning: if target_entropy: self.target_entropy = target_entropy else: self.target_entropy = -np.prod( self.env.action_space.shape).item( ) # heuristic value from Tuomas self.log_alpha = ptu.zeros(1, requires_grad=True) self.alpha_optimizer = optimizer_class( [self.log_alpha], lr=policy_lr, ) self.awr_use_mle_for_vf = awr_use_mle_for_vf self.vf_K = vf_K self.awr_sample_actions = awr_sample_actions self.awr_min_q = awr_min_q self.plotter = plotter self.render_eval_paths = render_eval_paths self.qf_criterion = nn.MSELoss() self.vf_criterion = nn.MSELoss() self.optimizers = {} self.policy_optimizer = optimizer_class( self.policy.parameters(), weight_decay=policy_weight_decay, lr=policy_lr, ) self.optimizers[self.policy] = self.policy_optimizer self.qf1_optimizer = optimizer_class( self.qf1.parameters(), weight_decay=q_weight_decay, lr=qf_lr, ) self.qf2_optimizer = optimizer_class( self.qf2.parameters(), weight_decay=q_weight_decay, lr=qf_lr, ) if buffer_policy and train_bc_on_rl_buffer: self.buffer_policy_optimizer = optimizer_class( self.buffer_policy.parameters(), weight_decay=policy_weight_decay, lr=policy_lr, ) self.optimizers[self.buffer_policy] = self.buffer_policy_optimizer self.optimizer_class = optimizer_class self.policy_weight_decay = policy_weight_decay self.policy_lr = policy_lr self.use_automatic_beta_tuning = use_automatic_beta_tuning and buffer_policy and train_bc_on_rl_buffer self.beta_epsilon = beta_epsilon if self.use_automatic_beta_tuning: self.log_beta = ptu.zeros(1, requires_grad=True) self.beta_optimizer = optimizer_class( [self.log_beta], lr=policy_lr, ) else: self.beta = beta self.beta_schedule_kwargs = beta_schedule_kwargs if beta_schedule_kwargs is None: self.beta_schedule = ConstantSchedule(beta) else: schedule_class = beta_schedule_kwargs.pop( "schedule_class", PiecewiseLinearSchedule) self.beta_schedule = schedule_class(**beta_schedule_kwargs) self.discount = discount self.reward_scale = reward_scale self.eval_statistics = OrderedDict() self._n_train_steps_total = 0 self._need_to_update_eval_statistics = True self.bc_num_pretrain_steps = bc_num_pretrain_steps self.q_num_pretrain1_steps = q_num_pretrain1_steps self.q_num_pretrain2_steps = q_num_pretrain2_steps self.bc_batch_size = bc_batch_size self.rl_weight = rl_weight self.bc_weight = bc_weight self.eval_policy = MakeDeterministic(self.policy) self.compute_bc = compute_bc self.alpha = alpha self.q_update_period = q_update_period self.policy_update_period = policy_update_period self.weight_loss = weight_loss self.reparam_weight = reparam_weight self.awr_weight = awr_weight self.post_pretrain_hyperparams = post_pretrain_hyperparams self.post_bc_pretrain_hyperparams = post_bc_pretrain_hyperparams self.update_policy = True self.pretraining_logging_period = pretraining_logging_period self.normalize_over_batch = normalize_over_batch self.normalize_over_state = normalize_over_state self.Z_K = Z_K self.reward_transform_class = reward_transform_class or LinearTransform self.reward_transform_kwargs = reward_transform_kwargs or dict(m=1, b=0) self.terminal_transform_class = terminal_transform_class or LinearTransform self.terminal_transform_kwargs = terminal_transform_kwargs or dict(m=1, b=0) self.reward_transform = self.reward_transform_class( **self.reward_transform_kwargs) self.terminal_transform = self.terminal_transform_class( **self.terminal_transform_kwargs) self.use_reparam_update = use_reparam_update self.clip_score = clip_score self.buffer_policy_sample_actions = buffer_policy_sample_actions self.train_bc_on_rl_buffer = train_bc_on_rl_buffer and buffer_policy self.validation_qlearning = validation_qlearning self.brac = brac self.mask_positive_advantage = mask_positive_advantage self.buffer_policy_reset_period = buffer_policy_reset_period self.num_buffer_policy_train_steps_on_reset = num_buffer_policy_train_steps_on_reset self.advantage_weighted_buffer_loss = advantage_weighted_buffer_loss
class AWACTrainer(TorchTrainer): def __init__( self, env, policy, qf1, qf2, target_qf1, target_qf2, buffer_policy=None, discount=0.99, reward_scale=1.0, beta=1.0, beta_schedule_kwargs=None, policy_lr=1e-3, qf_lr=1e-3, policy_weight_decay=0, q_weight_decay=0, optimizer_class=optim.Adam, soft_target_tau=1e-2, target_update_period=1, plotter=None, render_eval_paths=False, use_automatic_entropy_tuning=True, target_entropy=None, bc_num_pretrain_steps=0, q_num_pretrain1_steps=0, q_num_pretrain2_steps=0, bc_batch_size=128, alpha=1.0, policy_update_period=1, q_update_period=1, weight_loss=True, compute_bc=True, use_awr_update=True, use_reparam_update=False, bc_weight=0.0, rl_weight=1.0, reparam_weight=1.0, awr_weight=1.0, post_pretrain_hyperparams=None, post_bc_pretrain_hyperparams=None, awr_use_mle_for_vf=False, vf_K=1, awr_sample_actions=False, buffer_policy_sample_actions=False, awr_min_q=False, brac=False, reward_transform_class=None, reward_transform_kwargs=None, terminal_transform_class=None, terminal_transform_kwargs=None, pretraining_logging_period=1000, train_bc_on_rl_buffer=False, use_automatic_beta_tuning=False, beta_epsilon=1e-10, normalize_over_batch=True, normalize_over_state="advantage", Z_K=10, clip_score=None, validation_qlearning=False, mask_positive_advantage=False, buffer_policy_reset_period=-1, num_buffer_policy_train_steps_on_reset=100, advantage_weighted_buffer_loss=True, ): super().__init__() self.env = env self.policy = policy self.qf1 = qf1 self.qf2 = qf2 self.target_qf1 = target_qf1 self.target_qf2 = target_qf2 self.buffer_policy = buffer_policy self.soft_target_tau = soft_target_tau self.target_update_period = target_update_period self.use_awr_update = use_awr_update self.use_automatic_entropy_tuning = use_automatic_entropy_tuning if self.use_automatic_entropy_tuning: if target_entropy: self.target_entropy = target_entropy else: self.target_entropy = -np.prod( self.env.action_space.shape).item( ) # heuristic value from Tuomas self.log_alpha = ptu.zeros(1, requires_grad=True) self.alpha_optimizer = optimizer_class( [self.log_alpha], lr=policy_lr, ) self.awr_use_mle_for_vf = awr_use_mle_for_vf self.vf_K = vf_K self.awr_sample_actions = awr_sample_actions self.awr_min_q = awr_min_q self.plotter = plotter self.render_eval_paths = render_eval_paths self.qf_criterion = nn.MSELoss() self.vf_criterion = nn.MSELoss() self.optimizers = {} self.policy_optimizer = optimizer_class( self.policy.parameters(), weight_decay=policy_weight_decay, lr=policy_lr, ) self.optimizers[self.policy] = self.policy_optimizer self.qf1_optimizer = optimizer_class( self.qf1.parameters(), weight_decay=q_weight_decay, lr=qf_lr, ) self.qf2_optimizer = optimizer_class( self.qf2.parameters(), weight_decay=q_weight_decay, lr=qf_lr, ) if buffer_policy and train_bc_on_rl_buffer: self.buffer_policy_optimizer = optimizer_class( self.buffer_policy.parameters(), weight_decay=policy_weight_decay, lr=policy_lr, ) self.optimizers[self.buffer_policy] = self.buffer_policy_optimizer self.optimizer_class = optimizer_class self.policy_weight_decay = policy_weight_decay self.policy_lr = policy_lr self.use_automatic_beta_tuning = use_automatic_beta_tuning and buffer_policy and train_bc_on_rl_buffer self.beta_epsilon = beta_epsilon if self.use_automatic_beta_tuning: self.log_beta = ptu.zeros(1, requires_grad=True) self.beta_optimizer = optimizer_class( [self.log_beta], lr=policy_lr, ) else: self.beta = beta self.beta_schedule_kwargs = beta_schedule_kwargs if beta_schedule_kwargs is None: self.beta_schedule = ConstantSchedule(beta) else: schedule_class = beta_schedule_kwargs.pop( "schedule_class", PiecewiseLinearSchedule) self.beta_schedule = schedule_class(**beta_schedule_kwargs) self.discount = discount self.reward_scale = reward_scale self.eval_statistics = OrderedDict() self._n_train_steps_total = 0 self._need_to_update_eval_statistics = True self.bc_num_pretrain_steps = bc_num_pretrain_steps self.q_num_pretrain1_steps = q_num_pretrain1_steps self.q_num_pretrain2_steps = q_num_pretrain2_steps self.bc_batch_size = bc_batch_size self.rl_weight = rl_weight self.bc_weight = bc_weight self.eval_policy = MakeDeterministic(self.policy) self.compute_bc = compute_bc self.alpha = alpha self.q_update_period = q_update_period self.policy_update_period = policy_update_period self.weight_loss = weight_loss self.reparam_weight = reparam_weight self.awr_weight = awr_weight self.post_pretrain_hyperparams = post_pretrain_hyperparams self.post_bc_pretrain_hyperparams = post_bc_pretrain_hyperparams self.update_policy = True self.pretraining_logging_period = pretraining_logging_period self.normalize_over_batch = normalize_over_batch self.normalize_over_state = normalize_over_state self.Z_K = Z_K self.reward_transform_class = reward_transform_class or LinearTransform self.reward_transform_kwargs = reward_transform_kwargs or dict(m=1, b=0) self.terminal_transform_class = terminal_transform_class or LinearTransform self.terminal_transform_kwargs = terminal_transform_kwargs or dict(m=1, b=0) self.reward_transform = self.reward_transform_class( **self.reward_transform_kwargs) self.terminal_transform = self.terminal_transform_class( **self.terminal_transform_kwargs) self.use_reparam_update = use_reparam_update self.clip_score = clip_score self.buffer_policy_sample_actions = buffer_policy_sample_actions self.train_bc_on_rl_buffer = train_bc_on_rl_buffer and buffer_policy self.validation_qlearning = validation_qlearning self.brac = brac self.mask_positive_advantage = mask_positive_advantage self.buffer_policy_reset_period = buffer_policy_reset_period self.num_buffer_policy_train_steps_on_reset = num_buffer_policy_train_steps_on_reset self.advantage_weighted_buffer_loss = advantage_weighted_buffer_loss def get_batch_from_buffer(self, replay_buffer, batch_size): batch = replay_buffer.random_batch(batch_size) batch = np_to_pytorch_batch(batch) return batch def run_bc_batch(self, replay_buffer, policy): batch = self.get_batch_from_buffer(replay_buffer, self.bc_batch_size) o = batch["observations"] u = batch["actions"] # g = batch["resampled_goals"] # og = torch.cat((o, g), dim=1) og = o # pred_u, *_ = self.policy(og) dist = policy(og) pred_u, log_pi = dist.rsample_and_logprob() stats = dist.get_diagnostics() mse = (pred_u - u)**2 mse_loss = mse.mean() policy_logpp = dist.log_prob(u, ) logp_loss = -policy_logpp.mean() policy_loss = logp_loss return policy_loss, logp_loss, mse_loss, stats def pretrain_policy_with_bc( self, policy, train_buffer, test_buffer, steps, label="policy", ): logger.remove_tabular_output( 'progress.csv', relative_to_snapshot_dir=True, ) logger.add_tabular_output( 'pretrain_%s.csv' % label, relative_to_snapshot_dir=True, ) optimizer = self.optimizers[policy] prev_time = time.time() for i in range(steps): train_policy_loss, train_logp_loss, train_mse_loss, train_stats = self.run_bc_batch( train_buffer, policy) train_policy_loss = train_policy_loss * self.bc_weight optimizer.zero_grad() train_policy_loss.backward() optimizer.step() test_policy_loss, test_logp_loss, test_mse_loss, test_stats = self.run_bc_batch( test_buffer, policy) test_policy_loss = test_policy_loss * self.bc_weight if i % self.pretraining_logging_period == 0: stats = { "pretrain_bc/batch": i, "pretrain_bc/Train Logprob Loss": ptu.get_numpy(train_logp_loss), "pretrain_bc/Test Logprob Loss": ptu.get_numpy(test_logp_loss), "pretrain_bc/Train MSE": ptu.get_numpy(train_mse_loss), "pretrain_bc/Test MSE": ptu.get_numpy(test_mse_loss), "pretrain_bc/train_policy_loss": ptu.get_numpy(train_policy_loss), "pretrain_bc/test_policy_loss": ptu.get_numpy(test_policy_loss), "pretrain_bc/epoch_time": time.time() - prev_time, } logger.record_dict(stats) logger.dump_tabular(with_prefix=True, with_timestamp=False) pickle.dump( self.policy, open(logger.get_snapshot_dir() + '/bc_%s.pkl' % label, "wb")) prev_time = time.time() logger.remove_tabular_output( 'pretrain_%s.csv' % label, relative_to_snapshot_dir=True, ) logger.add_tabular_output( 'progress.csv', relative_to_snapshot_dir=True, ) if self.post_bc_pretrain_hyperparams: self.set_algorithm_weights(**self.post_bc_pretrain_hyperparams) def pretrain_q_with_bc_data(self): logger.remove_tabular_output('progress.csv', relative_to_snapshot_dir=True) logger.add_tabular_output('pretrain_q.csv', relative_to_snapshot_dir=True) self.update_policy = False # first train only the Q function for i in range(self.q_num_pretrain1_steps): self.eval_statistics = dict() train_data = self.replay_buffer.random_batch(self.bc_batch_size) train_data = np_to_pytorch_batch(train_data) obs = train_data['observations'] next_obs = train_data['next_observations'] # goals = train_data['resampled_goals'] train_data['observations'] = obs # torch.cat((obs, goals), dim=1) train_data[ 'next_observations'] = next_obs # torch.cat((next_obs, goals), dim=1) self.train_from_torch(train_data, pretrain=True) if i % self.pretraining_logging_period == 0: stats_with_prefix = add_prefix(self.eval_statistics, prefix="trainer/") logger.record_dict(stats_with_prefix) logger.dump_tabular(with_prefix=True, with_timestamp=False) self.update_policy = True # then train policy and Q function together prev_time = time.time() for i in range(self.q_num_pretrain2_steps): self.eval_statistics = dict() if i % self.pretraining_logging_period == 0: self._need_to_update_eval_statistics = True train_data = self.replay_buffer.random_batch(self.bc_batch_size) train_data = np_to_pytorch_batch(train_data) obs = train_data['observations'] next_obs = train_data['next_observations'] # goals = train_data['resampled_goals'] train_data['observations'] = obs # torch.cat((obs, goals), dim=1) train_data[ 'next_observations'] = next_obs # torch.cat((next_obs, goals), dim=1) self.train_from_torch(train_data, pretrain=True) if i % self.pretraining_logging_period == 0: self.eval_statistics["batch"] = i self.eval_statistics["epoch_time"] = time.time() - prev_time stats_with_prefix = add_prefix(self.eval_statistics, prefix="trainer/") logger.record_dict(stats_with_prefix) logger.dump_tabular(with_prefix=True, with_timestamp=False) prev_time = time.time() logger.remove_tabular_output( 'pretrain_q.csv', relative_to_snapshot_dir=True, ) logger.add_tabular_output( 'progress.csv', relative_to_snapshot_dir=True, ) self._need_to_update_eval_statistics = True self.eval_statistics = dict() if self.post_pretrain_hyperparams: self.set_algorithm_weights(**self.post_pretrain_hyperparams) def set_algorithm_weights(self, **kwargs): for key in kwargs: self.__dict__[key] = kwargs[key] def test_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 """ dist = self.policy(obs) new_obs_actions, log_pi = dist.rsample_and_logprob() policy_mle = dist.mle_estimate() if self.use_automatic_entropy_tuning: alpha_loss = -(self.log_alpha * (log_pi + self.target_entropy).detach()).mean() alpha = self.log_alpha.exp() else: alpha_loss = 0 alpha = self.alpha q1_pred = self.qf1(obs, actions) q2_pred = self.qf2(obs, actions) # Make sure policy accounts for squashing functions like tanh correctly! next_dist = self.policy(next_obs) new_next_actions, new_log_pi = next_dist.rsample_and_logprob() 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()) 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, ) policy_loss = (log_pi - q_new_actions).mean() self.eval_statistics['validation/QF1 Loss'] = np.mean( ptu.get_numpy(qf1_loss)) self.eval_statistics['validation/QF2 Loss'] = np.mean( ptu.get_numpy(qf2_loss)) self.eval_statistics['validation/Policy Loss'] = np.mean( ptu.get_numpy(policy_loss)) self.eval_statistics.update( create_stats_ordered_dict( 'validation/Q1 Predictions', ptu.get_numpy(q1_pred), )) self.eval_statistics.update( create_stats_ordered_dict( 'validation/Q2 Predictions', ptu.get_numpy(q2_pred), )) self.eval_statistics.update( create_stats_ordered_dict( 'validation/Q Targets', ptu.get_numpy(q_target), )) self.eval_statistics.update( create_stats_ordered_dict( 'validation/Log Pis', ptu.get_numpy(log_pi), )) policy_statistics = add_prefix(dist.get_diagnostics(), "validation/policy/") self.eval_statistics.update(policy_statistics) def train_from_torch( self, batch, train=True, pretrain=False, ): 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 """ dist = self.policy(obs) new_obs_actions, log_pi = dist.rsample_and_logprob() policy_mle = dist.mle_estimate() if self.brac: buf_dist = self.buffer_policy(obs) buf_log_pi = buf_dist.log_prob(actions) rewards = rewards + buf_log_pi 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! next_dist = self.policy(next_obs) new_next_actions, new_log_pi = next_dist.rsample_and_logprob() 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: v1_pi = self.qf1(obs, policy_mle) v2_pi = self.qf2(obs, policy_mle) v_pi = torch.min(v1_pi, v2_pi) else: if self.vf_K > 1: vs = [] for i in range(self.vf_K): u = dist.sample() q1 = self.qf1(obs, u) q2 = self.qf2(obs, u) v = torch.min(q1, q2) # v = q1 vs.append(v) v_pi = torch.cat(vs, 1).mean(dim=1) else: # v_pi = self.qf1(obs, new_obs_actions) v1_pi = self.qf1(obs, new_obs_actions) v2_pi = self.qf2(obs, new_obs_actions) v_pi = torch.min(v1_pi, v2_pi) 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 elif self.buffer_policy_sample_actions: buf_dist = self.buffer_policy(obs) u, _ = buf_dist.rsample_and_logprob() qf1_buffer_actions = self.qf1(obs, u) qf2_buffer_actions = self.qf2(obs, u) q_buffer_actions = torch.min( qf1_buffer_actions, qf2_buffer_actions, ) if self.awr_min_q: q_adv = q_buffer_actions else: q_adv = qf1_buffer_actions else: u = actions if self.awr_min_q: q_adv = torch.min(q1_pred, q2_pred) else: q_adv = q1_pred policy_logpp = dist.log_prob(u) if self.use_automatic_beta_tuning: buffer_dist = self.buffer_policy(obs) 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) if self.normalize_over_state == "advantage": score = q_adv - v_pi if self.mask_positive_advantage: score = torch.sign(score) elif self.normalize_over_state == "Z": buffer_dist = self.buffer_policy(obs) K = self.Z_K buffer_obs = [] buffer_actions = [] log_bs = [] log_pis = [] for i in range(K): u = buffer_dist.sample() log_b = buffer_dist.log_prob(u) log_pi = dist.log_prob(u) buffer_obs.append(obs) buffer_actions.append(u) log_bs.append(log_b) log_pis.append(log_pi) buffer_obs = torch.cat(buffer_obs, 0) buffer_actions = torch.cat(buffer_actions, 0) p_buffer = torch.exp(torch.cat(log_bs, 0).sum(dim=1, )) log_pi = torch.cat(log_pis, 0) log_pi = log_pi.sum(dim=1, ) q1_b = self.qf1(buffer_obs, buffer_actions) q2_b = self.qf2(buffer_obs, buffer_actions) q_b = torch.min(q1_b, q2_b) q_b = torch.reshape(q_b, (-1, K)) adv_b = q_b - v_pi # if self._n_train_steps_total % 100 == 0: # import ipdb; ipdb.set_trace() # Z = torch.exp(adv_b / beta).mean(dim=1, keepdim=True) # score = torch.exp((q_adv - v_pi) / beta) / Z # score = score / sum(score) logK = torch.log(ptu.tensor(float(K))) logZ = torch.logsumexp(adv_b / beta - logK, dim=1, keepdim=True) logS = (q_adv - v_pi) / beta - logZ # logZ = torch.logsumexp(q_b/beta - logK, dim=1, keepdim=True) # logS = q_adv/beta - logZ score = F.softmax(logS, dim=0) # score / sum(score) else: error if self.clip_score is not None: score = torch.clamp(score, max=self.clip_score) if self.weight_loss and weights is None: if self.normalize_over_batch == True: weights = F.softmax(score / beta, dim=0) elif self.normalize_over_batch == "whiten": adv_mean = torch.mean(score) adv_std = torch.std(score) + 1e-5 normalized_score = (score - adv_mean) / adv_std weights = torch.exp(normalized_score / beta) elif self.normalize_over_batch == "exp": weights = torch.exp(score / beta) elif self.normalize_over_batch == "step_fn": weights = (score > 0).float() elif self.normalize_over_batch == False: weights = score else: error weights = weights[:, 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 not pretrain and self.buffer_policy_reset_period > 0 and self._n_train_steps_total % self.buffer_policy_reset_period == 0: del self.buffer_policy_optimizer self.buffer_policy_optimizer = self.optimizer_class( self.buffer_policy.parameters(), weight_decay=self.policy_weight_decay, lr=self.policy_lr, ) self.optimizers[self.buffer_policy] = self.buffer_policy_optimizer for i in range(self.num_buffer_policy_train_steps_on_reset): if self.train_bc_on_rl_buffer: if self.advantage_weighted_buffer_loss: buffer_dist = self.buffer_policy(obs) buffer_u = actions buffer_new_obs_actions, _ = buffer_dist.rsample_and_logprob( ) buffer_policy_logpp = buffer_dist.log_prob(buffer_u) buffer_policy_logpp = buffer_policy_logpp[:, None] buffer_q1_pred = self.qf1(obs, buffer_u) buffer_q2_pred = self.qf2(obs, buffer_u) buffer_q_adv = torch.min(buffer_q1_pred, buffer_q2_pred) buffer_v1_pi = self.qf1(obs, buffer_new_obs_actions) buffer_v2_pi = self.qf2(obs, buffer_new_obs_actions) buffer_v_pi = torch.min(buffer_v1_pi, buffer_v2_pi) buffer_score = buffer_q_adv - buffer_v_pi buffer_weights = F.softmax(buffer_score / beta, dim=0) buffer_policy_loss = self.awr_weight * ( -buffer_policy_logpp * len(buffer_weights) * buffer_weights.detach()).mean() else: buffer_policy_loss, buffer_train_logp_loss, buffer_train_mse_loss, _ = self.run_bc_batch( self.replay_buffer.train_replay_buffer, self.buffer_policy) self.buffer_policy_optimizer.zero_grad() buffer_policy_loss.backward(retain_graph=True) self.buffer_policy_optimizer.step() if self.train_bc_on_rl_buffer: if self.advantage_weighted_buffer_loss: buffer_dist = self.buffer_policy(obs) buffer_u = actions buffer_new_obs_actions, _ = buffer_dist.rsample_and_logprob() buffer_policy_logpp = buffer_dist.log_prob(buffer_u) buffer_policy_logpp = buffer_policy_logpp[:, None] buffer_q1_pred = self.qf1(obs, buffer_u) buffer_q2_pred = self.qf2(obs, buffer_u) buffer_q_adv = torch.min(buffer_q1_pred, buffer_q2_pred) buffer_v1_pi = self.qf1(obs, buffer_new_obs_actions) buffer_v2_pi = self.qf2(obs, buffer_new_obs_actions) buffer_v_pi = torch.min(buffer_v1_pi, buffer_v2_pi) buffer_score = buffer_q_adv - buffer_v_pi buffer_weights = F.softmax(buffer_score / beta, dim=0) buffer_policy_loss = self.awr_weight * ( -buffer_policy_logpp * len(buffer_weights) * buffer_weights.detach()).mean() else: buffer_policy_loss, buffer_train_logp_loss, buffer_train_mse_loss, _ = self.run_bc_batch( self.replay_buffer.train_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( 'rewards', ptu.get_numpy(rewards), )) self.eval_statistics.update( create_stats_ordered_dict( 'terminals', ptu.get_numpy(terminals), )) policy_statistics = add_prefix(dist.get_diagnostics(), "policy/") self.eval_statistics.update(policy_statistics) self.eval_statistics.update( create_stats_ordered_dict( 'Advantage Weights', ptu.get_numpy(weights), )) self.eval_statistics.update( create_stats_ordered_dict( 'Advantage Score', ptu.get_numpy(score), )) if self.normalize_over_state == "Z": self.eval_statistics.update( create_stats_ordered_dict( 'logZ', ptu.get_numpy(logZ), )) 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: _, buffer_train_logp_loss, _, _ = self.run_bc_batch( self.replay_buffer.train_replay_buffer, self.buffer_policy) _, buffer_test_logp_loss, _, _ = self.run_bc_batch( self.replay_buffer.validation_replay_buffer, self.buffer_policy) buffer_dist = self.buffer_policy(obs) kldiv = torch.distributions.kl.kl_divergence(dist, buffer_dist) _, train_offline_logp_loss, _, _ = self.run_bc_batch( self.demo_train_buffer, self.buffer_policy) _, test_offline_logp_loss, _, _ = self.run_bc_batch( self.demo_test_buffer, self.buffer_policy) self.eval_statistics.update({ "buffer_policy/Train Online Logprob": -1 * ptu.get_numpy(buffer_train_logp_loss), "buffer_policy/Test Online Logprob": -1 * ptu.get_numpy(buffer_test_logp_loss), "buffer_policy/Train Offline Logprob": -1 * ptu.get_numpy(train_offline_logp_loss), "buffer_policy/Test Offline Logprob": -1 * ptu.get_numpy(test_offline_logp_loss), "buffer_policy/train_policy_loss": ptu.get_numpy(buffer_policy_loss), # "buffer_policy/test_policy_loss": ptu.get_numpy(buffer_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()), }) if self.validation_qlearning: train_data = self.replay_buffer.validation_replay_buffer.random_batch( self.bc_batch_size) train_data = np_to_pytorch_batch(train_data) obs = train_data['observations'] next_obs = train_data['next_observations'] # goals = train_data['resampled_goals'] train_data[ 'observations'] = obs # torch.cat((obs, goals), dim=1) train_data[ 'next_observations'] = next_obs # torch.cat((next_obs, goals), dim=1) self.test_from_torch(train_data) self._n_train_steps_total += 1 def get_diagnostics(self): stats = super().get_diagnostics() stats.update(self.eval_statistics) return stats def end_epoch(self, epoch): self._need_to_update_eval_statistics = True @property def networks(self): nets = [ self.policy, self.qf1, self.qf2, self.target_qf1, self.target_qf2, ] if self.buffer_policy: nets.append(self.buffer_policy) return nets def get_snapshot(self): return dict( policy=self.policy, qf1=self.qf1, qf2=self.qf2, target_qf1=self.qf1, target_qf2=self.qf2, buffer_policy=self.buffer_policy, )
def __init__( self, max_tau=10, epoch_max_tau_schedule=None, vectorized=True, cycle_taus_for_rollout=True, dense_rewards=False, finite_horizon=True, tau_sample_strategy='uniform', goal_reached_epsilon=1e-3, terminate_when_goal_reached=False, truncated_geom_factor=2., square_distance=False, goal_weights=None, normalize_distance=False, observation_key=None, desired_goal_key=None, ): """ :param max_tau: Maximum tau (planning horizon) to train with. :param epoch_max_tau_schedule: A schedule for the maximum planning horizon tau. :param vectorized: Train the QF in vectorized form? :param cycle_taus_for_rollout: Decrement the tau passed into the policy during rollout? :param dense_rewards: If True, always give rewards. Otherwise, only give rewards when the episode terminates. :param finite_horizon: If True, use a finite horizon formulation: give the time as input to the Q-function and terminate. :param tau_sample_strategy: Sampling strategy for taus used during training. Can be one of the following strings: - no_resampling: Do not resample the tau. Use the one from rollout. - uniform: Sample uniformly from [0, max_tau] - truncated_geometric: Sample from a truncated geometric distribution, truncated at max_tau. - all_valid: Always use all 0 to max_tau values :param goal_reached_epsilon: Epsilon used to determine if the goal has been reached. Used by `indicator` version of `reward_type` and when `terminate_whe_goal_reached` is True. :param terminate_when_goal_reached: Do you terminate when you have reached the goal? :param goal_weights: None or the weights for the different goal dimensions. These weights are used to compute the distances to the goal. """ assert tau_sample_strategy in [ 'no_resampling', 'uniform', 'truncated_geometric', 'all_valid', ] if epoch_max_tau_schedule is None: epoch_max_tau_schedule = ConstantSchedule(max_tau) if not finite_horizon: max_tau = 0 epoch_max_tau_schedule = ConstantSchedule(max_tau) cycle_taus_for_rollout = False self.max_tau = max_tau self.epoch_max_tau_schedule = epoch_max_tau_schedule self.vectorized = vectorized self.cycle_taus_for_rollout = cycle_taus_for_rollout self.dense_rewards = dense_rewards self.finite_horizon = finite_horizon self.tau_sample_strategy = tau_sample_strategy self.goal_reached_epsilon = goal_reached_epsilon self.terminate_when_goal_reached = terminate_when_goal_reached self.square_distance = square_distance self._rollout_tau = np.array([self.max_tau]) self.truncated_geom_factor = float(truncated_geom_factor) self.goal_weights = goal_weights if self.goal_weights is not None: # In case they were passed in as (e.g.) tuples or list self.goal_weights = np.array(self.goal_weights) assert self.goal_weights.size == self.env.goal_dim self.normalize_distance = normalize_distance self.observation_key = observation_key self.desired_goal_key = desired_goal_key self.eval_sampler = MultigoalSimplePathSampler( env=self.env, policy=self.eval_policy, max_samples=self.num_steps_per_eval, max_path_length=self.max_path_length, tau_sampling_function=self._sample_max_tau_for_rollout, cycle_taus_for_rollout=self.cycle_taus_for_rollout, render=self.render_during_eval, observation_key=self.observation_key, desired_goal_key=self.desired_goal_key, ) self.pretrain_obs = None
def __init__( self, model, batch_size=128, log_interval=0, beta=0.5, beta_schedule=None, lr=None, do_scatterplot=False, normalize=False, mse_weight=0.1, is_auto_encoder=False, background_subtract=False, linearity_weight=0.0, distance_weight=0.0, loss_weights=None, use_linear_dynamics=False, use_parallel_dataloading=False, train_data_workers=2, skew_dataset=False, skew_config=None, priority_function_kwargs=None, start_skew_epoch=0, weight_decay=0, key_to_reconstruct='observations', num_epochs=None, ): #TODO:steven fix pickling assert not use_parallel_dataloading, "Have to fix pickling the dataloaders first" if skew_config is None: skew_config = {} self.log_interval = log_interval self.batch_size = batch_size self.beta = beta if is_auto_encoder: self.beta = 0 if lr is None: if is_auto_encoder: lr = 1e-2 else: lr = 1e-3 self.beta_schedule = beta_schedule self.num_epochs = num_epochs if self.beta_schedule is None or is_auto_encoder: self.beta_schedule = ConstantSchedule(self.beta) self.imsize = model.imsize self.do_scatterplot = do_scatterplot model.to(ptu.device) self.model = model self.representation_size = model.representation_size self.input_channels = model.input_channels self.imlength = model.imlength self.lr = lr params = list(self.model.parameters()) self.optimizer = optim.Adam( params, lr=self.lr, weight_decay=weight_decay, ) self.key_to_reconstruct = key_to_reconstruct self.use_parallel_dataloading = use_parallel_dataloading self.train_data_workers = train_data_workers self.skew_dataset = skew_dataset self.skew_config = skew_config self.start_skew_epoch = start_skew_epoch if priority_function_kwargs is None: self.priority_function_kwargs = dict() else: self.priority_function_kwargs = priority_function_kwargs if use_parallel_dataloading: self.train_dataset_pt = ImageDataset(train_dataset, should_normalize=True) self.test_dataset_pt = ImageDataset(test_dataset, should_normalize=True) if self.skew_dataset: base_sampler = InfiniteWeightedRandomSampler( self.train_dataset, self._train_weights) else: base_sampler = InfiniteRandomSampler(self.train_dataset) self.train_dataloader = DataLoader( self.train_dataset_pt, sampler=InfiniteRandomSampler(self.train_dataset), batch_size=batch_size, drop_last=False, num_workers=train_data_workers, pin_memory=True, ) self.test_dataloader = DataLoader( self.test_dataset_pt, sampler=InfiniteRandomSampler(self.test_dataset), batch_size=batch_size, drop_last=False, num_workers=0, pin_memory=True, ) self.train_dataloader = iter(self.train_dataloader) self.test_dataloader = iter(self.test_dataloader) self.normalize = normalize self.mse_weight = mse_weight self.background_subtract = background_subtract if self.normalize or self.background_subtract: self.train_data_mean = np.mean(self.train_dataset, axis=0) self.train_data_mean = normalize_image( np.uint8(self.train_data_mean)) self.linearity_weight = linearity_weight self.distance_weight = distance_weight self.loss_weights = loss_weights self.use_linear_dynamics = use_linear_dynamics self._extra_stats_to_log = None # stateful tracking variables, reset every epoch self.eval_statistics = collections.defaultdict(list) self.eval_data = collections.defaultdict(list) self.num_batches = 0
class VAETrainer(LossFunction): def __init__( self, model, batch_size=128, log_interval=0, beta=0.5, beta_schedule=None, lr=None, do_scatterplot=False, normalize=False, mse_weight=0.1, is_auto_encoder=False, background_subtract=False, linearity_weight=0.0, distance_weight=0.0, loss_weights=None, use_linear_dynamics=False, use_parallel_dataloading=False, train_data_workers=2, skew_dataset=False, skew_config=None, priority_function_kwargs=None, start_skew_epoch=0, weight_decay=0, key_to_reconstruct='observations', num_epochs=None, ): #TODO:steven fix pickling assert not use_parallel_dataloading, "Have to fix pickling the dataloaders first" if skew_config is None: skew_config = {} self.log_interval = log_interval self.batch_size = batch_size self.beta = beta if is_auto_encoder: self.beta = 0 if lr is None: if is_auto_encoder: lr = 1e-2 else: lr = 1e-3 self.beta_schedule = beta_schedule self.num_epochs = num_epochs if self.beta_schedule is None or is_auto_encoder: self.beta_schedule = ConstantSchedule(self.beta) self.imsize = model.imsize self.do_scatterplot = do_scatterplot model.to(ptu.device) self.model = model self.representation_size = model.representation_size self.input_channels = model.input_channels self.imlength = model.imlength self.lr = lr params = list(self.model.parameters()) self.optimizer = optim.Adam( params, lr=self.lr, weight_decay=weight_decay, ) self.key_to_reconstruct = key_to_reconstruct self.use_parallel_dataloading = use_parallel_dataloading self.train_data_workers = train_data_workers self.skew_dataset = skew_dataset self.skew_config = skew_config self.start_skew_epoch = start_skew_epoch if priority_function_kwargs is None: self.priority_function_kwargs = dict() else: self.priority_function_kwargs = priority_function_kwargs if use_parallel_dataloading: self.train_dataset_pt = ImageDataset(train_dataset, should_normalize=True) self.test_dataset_pt = ImageDataset(test_dataset, should_normalize=True) if self.skew_dataset: base_sampler = InfiniteWeightedRandomSampler( self.train_dataset, self._train_weights) else: base_sampler = InfiniteRandomSampler(self.train_dataset) self.train_dataloader = DataLoader( self.train_dataset_pt, sampler=InfiniteRandomSampler(self.train_dataset), batch_size=batch_size, drop_last=False, num_workers=train_data_workers, pin_memory=True, ) self.test_dataloader = DataLoader( self.test_dataset_pt, sampler=InfiniteRandomSampler(self.test_dataset), batch_size=batch_size, drop_last=False, num_workers=0, pin_memory=True, ) self.train_dataloader = iter(self.train_dataloader) self.test_dataloader = iter(self.test_dataloader) self.normalize = normalize self.mse_weight = mse_weight self.background_subtract = background_subtract if self.normalize or self.background_subtract: self.train_data_mean = np.mean(self.train_dataset, axis=0) self.train_data_mean = normalize_image( np.uint8(self.train_data_mean)) self.linearity_weight = linearity_weight self.distance_weight = distance_weight self.loss_weights = loss_weights self.use_linear_dynamics = use_linear_dynamics self._extra_stats_to_log = None # stateful tracking variables, reset every epoch self.eval_statistics = collections.defaultdict(list) self.eval_data = collections.defaultdict(list) self.num_batches = 0 @property def log_dir(self): return logger.get_snapshot_dir() def get_dataset_stats(self, data): torch_input = ptu.from_numpy(normalize_image(data)) mus, log_vars = self.model.encode(torch_input) mus = ptu.get_numpy(mus) mean = np.mean(mus, axis=0) std = np.std(mus, axis=0) return mus, mean, std def _kl_np_to_np(self, np_imgs): torch_input = ptu.from_numpy(normalize_image(np_imgs)) mu, log_var = self.model.encode(torch_input) return ptu.get_numpy( -torch.sum(1 + log_var - mu.pow(2) - log_var.exp(), dim=1)) def _reconstruction_squared_error_np_to_np(self, np_imgs): torch_input = ptu.from_numpy(normalize_image(np_imgs)) recons, *_ = self.model(torch_input) error = torch_input - recons return ptu.get_numpy((error**2).sum(dim=1)) def set_vae(self, vae): self.model = vae self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr) def get_batch(self, test_data=False, epoch=None): if self.use_parallel_dataloading: if test_data: dataloader = self.test_dataloader else: dataloader = self.train_dataloader samples = next(dataloader).to(ptu.device) return samples dataset = self.test_dataset if test_data else self.train_dataset skew = False if epoch is not None: skew = (self.start_skew_epoch < epoch) if not test_data and self.skew_dataset and skew: probs = self._train_weights / np.sum(self._train_weights) ind = np.random.choice( len(probs), self.batch_size, p=probs, ) else: ind = np.random.randint(0, len(dataset), self.batch_size) samples = normalize_image(dataset[ind, :]) if self.normalize: samples = ((samples - self.train_data_mean) + 1) / 2 if self.background_subtract: samples = samples - self.train_data_mean return ptu.from_numpy(samples) def get_debug_batch(self, train=True): dataset = self.train_dataset if train else self.test_dataset X, Y = dataset ind = np.random.randint(0, Y.shape[0], self.batch_size) X = X[ind, :] Y = Y[ind, :] return ptu.from_numpy(X), ptu.from_numpy(Y) def train_epoch(self, epoch, dataset, batches=100): start_time = time.time() for b in range(batches): self.train_batch(epoch, dataset.random_batch(self.batch_size)) self.eval_statistics["train/epoch_duration"].append(time.time() - start_time) def test_epoch(self, epoch, dataset, batches=10): start_time = time.time() for b in range(batches): self.test_batch(epoch, dataset.random_batch(self.batch_size)) self.eval_statistics["test/epoch_duration"].append(time.time() - start_time) def compute_loss(self, batch, epoch=-1, test=False): prefix = "test/" if test else "train/" beta = float(self.beta_schedule.get_value(epoch)) obs = batch[self.key_to_reconstruct] reconstructions, obs_distribution_params, latent_distribution_params = self.model( obs) log_prob = self.model.logprob(obs, obs_distribution_params) kle = self.model.kl_divergence(latent_distribution_params) loss = -1 * log_prob + beta * kle self.eval_statistics['epoch'] = epoch self.eval_statistics['beta'] = beta self.eval_statistics[prefix + "losses"].append(loss.item()) self.eval_statistics[prefix + "log_probs"].append(log_prob.item()) self.eval_statistics[prefix + "kles"].append(kle.item()) self.eval_statistics["num_train_batches"].append(self.num_batches) encoder_mean = self.model.get_encoding_from_latent_distribution_params( latent_distribution_params) z_data = ptu.get_numpy(encoder_mean.cpu()) for i in range(len(z_data)): self.eval_data[prefix + "zs"].append(z_data[i, :]) self.eval_data[prefix + "last_batch"] = (obs, reconstructions) return loss def train_batch(self, epoch, batch): self.num_batches += 1 self.model.train() self.optimizer.zero_grad() loss = self.compute_loss(batch, epoch, False) loss.backward() self.optimizer.step() #self.scheduler.step() def test_batch( self, epoch, batch, ): self.model.eval() loss = self.compute_loss(batch, epoch, True) def end_epoch(self, epoch): self.eval_statistics = collections.defaultdict(list) self.test_last_batch = None def get_diagnostics(self): stats = OrderedDict() for k in sorted(self.eval_statistics.keys()): stats[k] = np.mean(self.eval_statistics[k]) return stats def dump_scatterplot(self, z, epoch): try: import matplotlib.pyplot as plt except ImportError: logger.log(__file__ + ": Unable to load matplotlib. Consider " "setting do_scatterplot to False") return dim_and_stds = [(i, np.std(z[:, i])) for i in range(z.shape[1])] dim_and_stds = sorted(dim_and_stds, key=lambda x: x[1]) dim1 = dim_and_stds[-1][0] dim2 = dim_and_stds[-2][0] plt.figure(figsize=(8, 8)) plt.scatter(z[:, dim1], z[:, dim2], marker='o', edgecolor='none') if self.model.dist_mu is not None: x1 = self.model.dist_mu[dim1:dim1 + 1] y1 = self.model.dist_mu[dim2:dim2 + 1] x2 = (self.model.dist_mu[dim1:dim1 + 1] + self.model.dist_std[dim1:dim1 + 1]) y2 = (self.model.dist_mu[dim2:dim2 + 1] + self.model.dist_std[dim2:dim2 + 1]) plt.plot([x1, x2], [y1, y2], color='k', linestyle='-', linewidth=2) axes = plt.gca() axes.set_xlim([-6, 6]) axes.set_ylim([-6, 6]) axes.set_title('dim {} vs dim {}'.format(dim1, dim2)) plt.grid(True) save_file = osp.join(self.log_dir, 'scatter%d.png' % epoch) plt.savefig(save_file)
def train_vae(variant, return_data=False): from rlkit.misc.ml_util import PiecewiseLinearSchedule, ConstantSchedule from rlkit.torch.vae.conv_vae import ( ConvVAE, ConvDynamicsVAE, SpatialAutoEncoder, AutoEncoder, ) import rlkit.torch.vae.conv_vae as conv_vae from rlkit.torch.vae.vae_trainer import ConvVAETrainer from rlkit.core import logger import rlkit.torch.pytorch_util as ptu from rlkit.pythonplusplus import identity import torch import gym beta = variant["beta"] representation_size = variant.get( "representation_size", variant.get("latent_sizes", variant.get("embedding_dim", None))) use_linear_dynamics = variant.get('use_linear_dynamics', False) variant['algo_kwargs']['num_epochs'] = variant['num_epochs'] generate_vae_dataset_fctn = variant.get('generate_vae_data_fctn', generate_vae_dataset) variant['generate_vae_dataset_kwargs'][ 'use_linear_dynamics'] = use_linear_dynamics variant['generate_vae_dataset_kwargs']['batch_size'] = variant[ 'algo_kwargs']['batch_size'] train_dataset, test_dataset, info = generate_vae_dataset_fctn( variant['generate_vae_dataset_kwargs']) if use_linear_dynamics: action_dim = train_dataset.data['actions'].shape[2] logger.save_extra_data(info) logger.get_snapshot_dir() if 'beta_schedule_kwargs' in variant: beta_schedule = PiecewiseLinearSchedule( **variant['beta_schedule_kwargs']) else: beta_schedule = None if 'context_schedule' in variant: schedule = variant['context_schedule'] if type(schedule) is dict: context_schedule = PiecewiseLinearSchedule(**schedule) else: context_schedule = ConstantSchedule(schedule) variant['algo_kwargs']['context_schedule'] = context_schedule if variant.get('decoder_activation', None) == 'sigmoid': decoder_activation = torch.nn.Sigmoid() else: decoder_activation = identity architecture = variant['vae_kwargs'].get('architecture', None) if not architecture and variant.get('imsize') == 84: architecture = conv_vae.imsize84_default_architecture elif not architecture and variant.get('imsize') == 48: architecture = conv_vae.imsize48_default_architecture variant['vae_kwargs']['architecture'] = architecture variant['vae_kwargs']['imsize'] = variant.get('imsize') if variant['algo_kwargs'].get('is_auto_encoder', False): model = AutoEncoder(representation_size, decoder_output_activation=decoder_activation, **variant['vae_kwargs']) elif variant.get('use_spatial_auto_encoder', False): model = SpatialAutoEncoder( representation_size, decoder_output_activation=decoder_activation, **variant['vae_kwargs']) elif variant.get('only_kwargs', False): vae_class = variant.get('vae_class', ConvVAE) model = vae_class(**variant['vae_kwargs']) else: vae_class = variant.get('vae_class', ConvVAE) if use_linear_dynamics: model = vae_class(representation_size, decoder_output_activation=decoder_activation, action_dim=action_dim, **variant['vae_kwargs']) else: model = vae_class(representation_size, decoder_output_activation=decoder_activation, **variant['vae_kwargs']) model.to(ptu.device) vae_trainer_class = variant.get('vae_trainer_class', ConvVAETrainer) trainer = vae_trainer_class(model, beta=beta, beta_schedule=beta_schedule, **variant['algo_kwargs']) save_period = variant['save_period'] dump_skew_debug_plots = variant.get('dump_skew_debug_plots', False) for epoch in range(variant['num_epochs']): should_save_imgs = (epoch % save_period == 0) trainer.train_epoch(epoch, train_dataset) trainer.test_epoch(epoch, test_dataset) if should_save_imgs: trainer.dump_reconstructions(epoch) trainer.dump_samples(epoch) if dump_skew_debug_plots: trainer.dump_best_reconstruction(epoch) trainer.dump_worst_reconstruction(epoch) trainer.dump_sampling_histogram(epoch) stats = trainer.get_diagnostics() for k, v in stats.items(): logger.record_tabular(k, v) logger.dump_tabular() trainer.end_epoch(epoch) if epoch % 50 == 0: logger.save_itr_params(epoch, model) logger.save_extra_data(model, 'vae.pkl', mode='pickle') if return_data: return model, train_dataset, test_dataset return model
def __init__(self, env, qf, policy, target_qf, target_policy, exploration_policy, policy_learning_rate=1e-4, qf_learning_rate=1e-3, qf_weight_decay=0, target_hard_update_period=1000, tau=1e-2, use_soft_update=True, qf_criterion=None, residual_gradient_weight=0, epoch_discount_schedule=None, eval_with_target_policy=False, policy_pre_activation_weight=0., optimizer_class=optim.Adam, plotter=None, render_eval_paths=False, obs_normalizer: TorchFixedNormalizer = None, action_normalizer: TorchFixedNormalizer = None, num_paths_for_normalization=0, min_q_value=-np.inf, max_q_value=np.inf, **kwargs): """ :param env: :param qf: :param policy: :param exploration_policy: :param policy_learning_rate: :param qf_learning_rate: :param qf_weight_decay: :param target_hard_update_period: :param tau: :param use_soft_update: :param qf_criterion: Loss function to use for the q function. Should be a function that takes in two inputs (y_predicted, y_target). :param residual_gradient_weight: c, float between 0 and 1. The gradient used for training the Q function is then (1-c) * normal td gradient + c * residual gradient :param epoch_discount_schedule: A schedule for the discount factor that varies with the epoch. :param kwargs: """ self.target_policy = target_policy if eval_with_target_policy: eval_policy = self.target_policy else: eval_policy = policy super().__init__(env, exploration_policy, eval_policy=eval_policy, **kwargs) if qf_criterion is None: qf_criterion = nn.MSELoss() self.qf = qf self.policy = policy self.policy_learning_rate = policy_learning_rate self.qf_learning_rate = qf_learning_rate self.qf_weight_decay = qf_weight_decay self.target_hard_update_period = target_hard_update_period self.tau = tau self.use_soft_update = use_soft_update self.residual_gradient_weight = residual_gradient_weight self.policy_pre_activation_weight = policy_pre_activation_weight self.qf_criterion = qf_criterion if epoch_discount_schedule is None: epoch_discount_schedule = ConstantSchedule(self.discount) self.epoch_discount_schedule = epoch_discount_schedule self.plotter = plotter self.render_eval_paths = render_eval_paths self.obs_normalizer = obs_normalizer self.action_normalizer = action_normalizer self.num_paths_for_normalization = num_paths_for_normalization self.min_q_value = min_q_value self.max_q_value = max_q_value self.target_qf = target_qf self.qf_optimizer = optimizer_class( self.qf.parameters(), lr=self.qf_learning_rate, ) self.policy_optimizer = optimizer_class( self.policy.parameters(), lr=self.policy_learning_rate, )
def __init__( self, vae, *args, decoded_obs_key='image_observation', decoded_achieved_goal_key='image_achieved_goal', decoded_desired_goal_key='image_desired_goal', exploration_rewards_type='None', exploration_rewards_scale=1.0, vae_priority_type='None', start_skew_epoch=0, power=1.0, internal_keys=None, exploration_schedule_kwargs=None, priority_function_kwargs=None, exploration_counter_kwargs=None, relabeling_goal_sampling_mode='vae_prior', decode_vae_goals=False, **kwargs ): if internal_keys is None: internal_keys = [] for key in [ decoded_obs_key, decoded_achieved_goal_key, decoded_desired_goal_key ]: if key not in internal_keys: internal_keys.append(key) super().__init__(internal_keys=internal_keys, *args, **kwargs) # assert isinstance(self.env, VAEWrappedEnv) self.vae = vae self.decoded_obs_key = decoded_obs_key self.decoded_desired_goal_key = decoded_desired_goal_key self.decoded_achieved_goal_key = decoded_achieved_goal_key self.exploration_rewards_type = exploration_rewards_type self.exploration_rewards_scale = exploration_rewards_scale self.start_skew_epoch = start_skew_epoch self.vae_priority_type = vae_priority_type self.power = power self._relabeling_goal_sampling_mode = relabeling_goal_sampling_mode self.decode_vae_goals = decode_vae_goals if exploration_schedule_kwargs is None: self.explr_reward_scale_schedule = \ ConstantSchedule(self.exploration_rewards_scale) else: self.explr_reward_scale_schedule = \ PiecewiseLinearSchedule(**exploration_schedule_kwargs) self._give_explr_reward_bonus = ( exploration_rewards_type != 'None' and exploration_rewards_scale != 0. ) self._exploration_rewards = np.zeros((self.max_size, 1), dtype=np.float32) self._prioritize_vae_samples = ( vae_priority_type != 'None' and power != 0. ) self._vae_sample_priorities = np.zeros((self.max_size, 1), dtype=np.float32) self._vae_sample_probs = None self.use_dynamics_model = ( self.exploration_rewards_type == 'forward_model_error' ) if self.use_dynamics_model: self.initialize_dynamics_model() type_to_function = { 'reconstruction_error': self.reconstruction_mse, 'bce': self.binary_cross_entropy, 'latent_distance': self.latent_novelty, 'latent_distance_true_prior': self.latent_novelty_true_prior, 'forward_model_error': self.forward_model_error, 'gaussian_inv_prob': self.gaussian_inv_prob, 'bernoulli_inv_prob': self.bernoulli_inv_prob, 'vae_prob': self.vae_prob, 'hash_count': self.hash_count_reward, 'None': self.no_reward, } self.exploration_reward_func = ( type_to_function[self.exploration_rewards_type] ) self.vae_prioritization_func = ( type_to_function[self.vae_priority_type] ) if priority_function_kwargs is None: self.priority_function_kwargs = dict() else: self.priority_function_kwargs = priority_function_kwargs if self.exploration_rewards_type == 'hash_count': if exploration_counter_kwargs is None: exploration_counter_kwargs = dict() self.exploration_counter = CountExploration(env=self.env, **exploration_counter_kwargs) self.epoch = 0
def __init__(self, env, exploration_policy, beta_q, beta_q2, beta_v, policy, train_with='both', goal_reached_epsilon=1e-3, learning_rate=1e-3, prioritized_replay=False, always_reset_env=True, finite_horizon=False, max_num_steps_left=0, flip_training_period=100, train_simultaneously=True, policy_and_target_update_period=2, target_policy_noise=0.2, target_policy_noise_clip=0.5, soft_target_tau=0.005, per_beta_schedule=None, **kwargs): self.train_simultaneously = train_simultaneously assert train_with in ['both', 'off_policy', 'on_policy'] super().__init__(env, exploration_policy, **kwargs) self.eval_sampler = MultigoalSimplePathSampler( env=self.env, policy=self.eval_policy, max_samples=self.num_steps_per_eval, max_path_length=self.max_path_length, tau_sampling_function=lambda: 0, goal_sampling_function=self.env.sample_goal_for_rollout, cycle_taus_for_rollout=False, render=self.render_during_eval) self.goal_reached_epsilon = goal_reached_epsilon self.beta_q = beta_q self.beta_v = beta_v self.beta_q2 = beta_q2 self.target_beta_q = self.beta_q.copy() self.target_beta_q2 = self.beta_q2.copy() self.train_with = train_with self.policy = policy self.target_policy = policy self.prioritized_replay = prioritized_replay self.flip_training_period = flip_training_period self.always_reset_env = always_reset_env self.finite_horizon = finite_horizon self.max_num_steps_left = max_num_steps_left assert max_num_steps_left >= 0 self.policy_and_target_update_period = policy_and_target_update_period self.target_policy_noise = target_policy_noise self.target_policy_noise_clip = target_policy_noise_clip self.soft_target_tau = soft_target_tau if per_beta_schedule is None: per_beta_schedule = ConstantSchedule(1.0) self.per_beta_schedule = per_beta_schedule self.beta_q_optimizer = Adam(self.beta_q.parameters(), lr=learning_rate) self.beta_q2_optimizer = Adam(self.beta_q2.parameters(), lr=learning_rate) self.beta_v_optimizer = Adam(self.beta_v.parameters(), lr=learning_rate) self.policy_optimizer = Adam( self.policy.parameters(), lr=learning_rate, ) self.q_criterion = nn.BCELoss() self.v_criterion = nn.BCELoss() # For the multitask env self._rollout_goal = None self.extra_eval_statistics = OrderedDict() for key_not_always_updated in [ 'Policy Gradient Norms', 'Beta Q Gradient Norms', 'dQ/da', ]: self.eval_statistics.update( create_stats_ordered_dict( key_not_always_updated, np.zeros(2), )) self.training_policy = False # For debugging self.train_batches = []
def __init__( self, model, log_interval=0, beta=0.5, beta_schedule=None, lr=1e-3, do_scatterplot=False, normalize=False, mse_weight=0.1, is_auto_encoder=False, background_subtract=False, linearity_weight=0.0, distance_weight=0.0, loss_weights=None, use_linear_dynamics=False, use_parallel_dataloading=False, train_data_workers=2, skew_dataset=False, skew_config=None, priority_function_kwargs=None, start_skew_epoch=0, weight_decay=0, batch_size=64, ): #TODO:steven fix pickling assert not use_parallel_dataloading, "Have to fix pickling the dataloaders first" if skew_config is None: skew_config = {} self.log_interval = log_interval self.beta = beta if is_auto_encoder: self.beta = 0 self.beta_schedule = beta_schedule if self.beta_schedule is None or is_auto_encoder: self.beta_schedule = ConstantSchedule(self.beta) self.do_scatterplot = do_scatterplot model.to(ptu.device) self.model = model self.lr = lr params = list(self.model.parameters()) self.optimizer = optim.Adam( params, lr=self.lr, weight_decay=weight_decay, ) self.batch_size = batch_size self.use_parallel_dataloading = use_parallel_dataloading self.train_data_workers = train_data_workers self.skew_dataset = skew_dataset self.skew_config = skew_config self.start_skew_epoch = start_skew_epoch if priority_function_kwargs is None: self.priority_function_kwargs = dict() else: self.priority_function_kwargs = priority_function_kwargs self.normalize = normalize self.mse_weight = mse_weight self.background_subtract = background_subtract self.linearity_weight = linearity_weight self.distance_weight = distance_weight self.loss_weights = loss_weights self.loss_fn = torch.nn.CrossEntropyLoss() self.log_softmax = torch.nn.LogSoftmax() self.use_linear_dynamics = use_linear_dynamics self._extra_stats_to_log = None # stateful tracking variables, reset every epoch self.eval_statistics = collections.defaultdict(list) self.eval_data = collections.defaultdict(list) self.num_train_batches = 0 self.num_test_batches = 0 self.bin_midpoints = ( torch.arange(0, self.model.output_classes).float() + 0.5) / self.model.output_classes self.bin_midpoints = self.bin_midpoints.to(ptu.device)