class VAETrainer():
def __init__(
self,
train_dataset,
test_dataset,
model,
batch_size=128,
beta=0.5,
beta_schedule=None,
lr=1e-3,
extra_recon_logging=dict(),
recon_weights=None,
recon_loss_type='mse',
**kwargs
):
assert recon_loss_type in ['mse', 'wse']
self.batch_size = batch_size
self.beta = beta
self.beta_schedule = beta_schedule
if self.beta_schedule is None:
self.beta_schedule = ConstantSchedule(self.beta)
if ptu.gpu_enabled():
model.cuda()
self.model = model
self.representation_size = model.representation_size
self.optimizer = optim.Adam(self.model.parameters(), lr=lr)
self.train_dataset, self.test_dataset = train_dataset, test_dataset
assert self.train_dataset['next_obs'].dtype == np.float32
assert self.test_dataset['next_obs'].dtype ==np.float32
assert self.train_dataset['obs'].dtype == np.float32
assert self.test_dataset['obs'].dtype == np.float32
self.normalize = model.normalize
self.mse = nn.MSELoss()
if self.normalize:
self.train_data_mean = ptu.np_to_var(np.mean(self.train_dataset['next_obs'], axis=0))
np_std = np.std(self.train_dataset['next_obs'], axis=0)
for i in range(len(np_std)):
if np_std[i] < 1e-3:
np_std[i] = 1.0
self.train_data_std = ptu.np_to_var(np_std)
self.model.train_data_mean = self.train_data_mean
self.model.train_data_std = self.train_data_std
self.extra_recon_logging = extra_recon_logging
self.recon_weights = recon_weights
self.recon_loss_type = recon_loss_type
def get_batch(self, train=True):
dataset = self.train_dataset if train else self.test_dataset
ind = np.random.randint(0, len(dataset['obs']), self.batch_size)
samples_obs = dataset['obs'][ind, :]
samples_actions = dataset['actions'][ind, :]
samples_next_obs = dataset['next_obs'][ind, :]
return {
'obs': ptu.np_to_var(samples_obs),
'actions': ptu.np_to_var(samples_actions),
'next_obs': ptu.np_to_var(samples_next_obs),
}
def logprob(self, recon_x, x, normalize=None, idx=None, unorm_weights=None):
if normalize is None:
normalize = self.normalize
if normalize:
x = (x - self.train_data_mean) / self.train_data_std
recon_x = (recon_x - self.train_data_mean) / self.train_data_std
if idx is not None:
x = x[:,idx]
recon_x = recon_x[:,idx]
if unorm_weights is not None:
unorm_weights = unorm_weights[idx]
if unorm_weights is not None:
dim = x.shape[1]
norm_weights = unorm_weights / (np.sum(unorm_weights) / dim)
norm_weights = ptu.np_to_var(norm_weights)
recon_x = recon_x * norm_weights
x = x * norm_weights
return self.mse(recon_x, x)
def kl_divergence(self, mu, logvar):
kl = - torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1).mean()
return kl
def train_epoch(self, epoch, batches=100):
self.model.train()
losses = []
kles = []
mses = []
beta = self.beta_schedule.get_value(epoch)
for batch_idx in range(batches):
data = self.get_batch()
obs = data['obs']
next_obs = data['next_obs']
actions = data['actions']
self.optimizer.zero_grad()
recon_batch, mu, logvar = self.model(next_obs)
mse = self.logprob(recon_batch, next_obs)
kle = self.kl_divergence(mu, logvar)
if self.recon_loss_type == 'mse':
loss = mse + beta * kle
elif self.recon_loss_type == 'wse':
wse = self.logprob(recon_batch, next_obs, unorm_weights=self.recon_weights)
loss = wse + beta * kle
loss.backward()
losses.append(loss.data[0])
mses.append(mse.data[0])
kles.append(kle.data[0])
self.optimizer.step()
logger.record_tabular("train/epoch", epoch)
logger.record_tabular("train/MSE", np.mean(mses))
logger.record_tabular("train/KL", np.mean(kles))
logger.record_tabular("train/loss", np.mean(losses))
def test_epoch(self, epoch, save_vae=True, **kwargs):
self.model.eval()
losses = []
kles = []
zs = []
recon_logging_dict = {
'MSE': [],
'WSE': [],
}
for k in self.extra_recon_logging:
recon_logging_dict[k] = []
beta = self.beta_schedule.get_value(epoch)
for batch_idx in range(100):
data = self.get_batch(train=False)
obs = data['obs']
next_obs = data['next_obs']
actions = data['actions']
recon_batch, mu, logvar = self.model(next_obs)
mse = self.logprob(recon_batch, next_obs)
wse = self.logprob(recon_batch, next_obs, unorm_weights=self.recon_weights)
for k, idx in self.extra_recon_logging.items():
recon_loss = self.logprob(recon_batch, next_obs, idx=idx)
recon_logging_dict[k].append(recon_loss.data[0])
kle = self.kl_divergence(mu, logvar)
if self.recon_loss_type == 'mse':
loss = mse + beta * kle
elif self.recon_loss_type == 'wse':
loss = wse + beta * kle
z_data = ptu.get_numpy(mu.cpu())
for i in range(len(z_data)):
zs.append(z_data[i, :])
losses.append(loss.data[0])
recon_logging_dict['WSE'].append(wse.data[0])
recon_logging_dict['MSE'].append(mse.data[0])
kles.append(kle.data[0])
zs = np.array(zs)
self.model.dist_mu = zs.mean(axis=0)
self.model.dist_std = zs.std(axis=0)
for k in recon_logging_dict:
logger.record_tabular("/".join(["test", k]), np.mean(recon_logging_dict[k]))
logger.record_tabular("test/KL", np.mean(kles))
logger.record_tabular("test/loss", np.mean(losses))
logger.record_tabular("beta", beta)
process = psutil.Process(os.getpid())
logger.record_tabular("RAM Usage (Mb)", int(process.memory_info().rss / 1000000))
num_active_dims = 0
for std in self.model.dist_std:
if std > 0.15:
num_active_dims += 1
logger.record_tabular("num_active_dims", num_active_dims)
logger.dump_tabular()
if save_vae:
logger.save_itr_params(epoch, self.model, prefix='vae', save_anyway=True) # slow...
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",
):
#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
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.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
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)
@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())
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.model.train()
self.optimizer.zero_grad()
loss = self.compute_loss(batch, epoch, False)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.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)
class AWRSACTrainer(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,
bc_loss_type="mle",
awr_loss_type="mle",
save_bc_policies=0,
alpha=1.0,
policy_update_period=1,
q_update_period=1,
weight_loss=True,
compute_bc=True,
bc_weight=0.0,
rl_weight=1.0,
use_awr_update=True,
use_reparam_update=False,
reparam_weight=1.0,
awr_weight=1.0,
post_pretrain_hyperparams=None,
post_bc_pretrain_hyperparams=None,
awr_use_mle_for_vf=False,
awr_sample_actions=False,
awr_min_q=False,
reward_transform_class=None,
reward_transform_kwargs=None,
terminal_transform_class=None,
terminal_transform_kwargs=None,
pretraining_env_logging_period=100000,
pretraining_logging_period=1000,
do_pretrain_rollouts=False,
train_bc_on_rl_buffer=False,
use_automatic_beta_tuning=False,
beta_epsilon=1e-10,
):
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.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.policy_optimizer = optimizer_class(
self.policy.parameters(),
weight_decay=policy_weight_decay,
lr=policy_lr,
)
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.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.bc_loss_type = bc_loss_type
self.awr_loss_type = awr_loss_type
self.rl_weight = rl_weight
self.bc_weight = bc_weight
self.save_bc_policies = save_bc_policies
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_env_logging_period = pretraining_env_logging_period
self.pretraining_logging_period = pretraining_logging_period
self.do_pretrain_rollouts = do_pretrain_rollouts
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.train_bc_on_rl_buffer = train_bc_on_rl_buffer and buffer_policy
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)
pred_u, policy_mean, policy_log_std, log_pi, entropy, policy_std, mean_action_log_prob, pretanh_value, dist = policy(
og, deterministic=False, reparameterize=True, return_log_prob=True,
)
mse = (policy_mean - u) ** 2
mse_loss = mse.mean()
policy_logpp = dist.log_prob(u, )
logp_loss = -policy_logpp.mean()
# T = 0
if self.bc_loss_type == "mle":
policy_loss = logp_loss
elif self.bc_loss_type == "mse":
policy_loss = mse_loss
else:
error
return policy_loss, logp_loss, mse_loss, policy_log_std
def do_rollouts(self):
total_ret = 0
for _ in range(20):
o = self.env.reset()
ret = 0
for _ in range(1000):
a, _ = self.policy.get_action(o)
o, r, done, info = self.env.step(a)
ret += r
if done:
break
total_ret += ret
return total_ret
def pretrain_policy_with_bc(self):
logger.remove_tabular_output(
'progress.csv', relative_to_snapshot_dir=True
)
logger.add_tabular_output(
'pretrain_policy.csv', relative_to_snapshot_dir=True
)
if self.do_pretrain_rollouts:
total_ret = self.do_rollouts()
print("INITIAL RETURN", total_ret/20)
prev_time = time.time()
for i in range(self.bc_num_pretrain_steps):
train_policy_loss, train_logp_loss, train_mse_loss, train_log_std = self.run_bc_batch(self.demo_train_buffer, self.policy)
train_policy_loss = train_policy_loss * self.bc_weight
self.policy_optimizer.zero_grad()
train_policy_loss.backward()
self.policy_optimizer.step()
test_policy_loss, test_logp_loss, test_mse_loss, test_log_std = self.run_bc_batch(self.demo_test_buffer, self.policy)
test_policy_loss = test_policy_loss * self.bc_weight
if self.do_pretrain_rollouts and i % self.pretraining_env_logging_period == 0:
total_ret = self.do_rollouts()
print("Return at step {} : {}".format(i, total_ret/20))
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,
}
if self.do_pretrain_rollouts:
stats["pretrain_bc/avg_return"] = total_ret / 20
logger.record_dict(stats)
logger.dump_tabular(with_prefix=True, with_timestamp=False)
pickle.dump(self.policy, open(logger.get_snapshot_dir() + '/bc.pkl', "wb"))
prev_time = time.time()
logger.remove_tabular_output(
'pretrain_policy.csv',
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)
if i%self.pretraining_logging_period == 0:
logger.record_dict(self.eval_statistics)
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)
if self.do_pretrain_rollouts and i % self.pretraining_env_logging_period == 0:
total_ret = self.do_rollouts()
print("Return at step {} : {}".format(i, total_ret/20))
if i%self.pretraining_logging_period==0:
if self.do_pretrain_rollouts:
self.eval_statistics["pretrain_bc/avg_return"] = total_ret / 20
self.eval_statistics["batch"] = i
self.eval_statistics["epoch_time"] = time.time()-prev_time
logger.record_dict(self.eval_statistics)
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,
# bc_weight,
# rl_weight,
# use_awr_update,
# use_reparam_update,
# reparam_weight,
# awr_weight,
**kwargs
):
for key in kwargs:
self.__dict__[key] = kwargs[key]
# self.bc_weight = bc_weight
# self.rl_weight = rl_weight
# self.use_awr_update = use_awr_update
# self.use_reparam_update = use_reparam_update
# self.awr_weight = awr_weight
def train_from_torch(self, batch):
rewards = batch['rewards']
terminals = batch['terminals']
obs = batch['observations']
actions = batch['actions']
next_obs = batch['next_observations']
weights = batch.get('weights', None)
if self.reward_transform:
rewards = self.reward_transform(rewards)
if self.terminal_transform:
terminals = self.terminal_transform(terminals)
"""
Policy and Alpha Loss
"""
new_obs_actions, policy_mean, policy_log_std, log_pi, entropy, policy_std, mean_action_log_prob, pretanh_value, dist = self.policy(
obs, reparameterize=True, return_log_prob=True,
)
if self.use_automatic_entropy_tuning:
alpha_loss = -(self.log_alpha * (log_pi + self.target_entropy).detach()).mean()
self.alpha_optimizer.zero_grad()
alpha_loss.backward()
self.alpha_optimizer.step()
alpha = self.log_alpha.exp()
else:
alpha_loss = 0
alpha = self.alpha
"""
QF Loss
"""
q1_pred = self.qf1(obs, actions)
q2_pred = self.qf2(obs, actions)
# Make sure policy accounts for squashing functions like tanh correctly!
new_next_actions, _, _, new_log_pi, *_ = self.policy(
next_obs, reparameterize=True, return_log_prob=True,
)
target_q_values = torch.min(
self.target_qf1(next_obs, new_next_actions),
self.target_qf2(next_obs, new_next_actions),
) - alpha * new_log_pi
q_target = self.reward_scale * rewards + (1. - terminals) * self.discount * target_q_values
qf1_loss = self.qf_criterion(q1_pred, q_target.detach())
qf2_loss = self.qf_criterion(q2_pred, q_target.detach())
"""
Policy Loss
"""
qf1_new_actions = self.qf1(obs, new_obs_actions)
qf2_new_actions = self.qf2(obs, new_obs_actions)
q_new_actions = torch.min(
qf1_new_actions,
qf2_new_actions,
)
# Advantage-weighted regression
if self.awr_use_mle_for_vf:
v_pi = self.qf1(obs, policy_mean)
else:
v_pi = self.qf1(obs, new_obs_actions)
if self.awr_sample_actions:
u = new_obs_actions
if self.awr_min_q:
q_adv = q_new_actions
else:
q_adv = qf1_new_actions
else:
u = actions
if self.awr_min_q:
q_adv = torch.min(q1_pred, q2_pred)
else:
q_adv = q1_pred
if self.awr_loss_type == "mse":
policy_logpp = -(policy_mean - actions) ** 2
else:
policy_logpp = dist.log_prob(u)
policy_logpp = policy_logpp.sum(dim=1, keepdim=True)
advantage = q_adv - v_pi
if self.weight_loss and weights is None:
if self.use_automatic_beta_tuning:
_, _, _, _, _, _, _, _, buffer_dist = self.buffer_policy(
obs, reparameterize=True, return_log_prob=True,
)
beta = self.log_beta.exp()
kldiv = torch.distributions.kl.kl_divergence(dist, buffer_dist)
beta_loss = -1*(beta*(kldiv-self.beta_epsilon).detach()).mean()
self.beta_optimizer.zero_grad()
beta_loss.backward()
self.beta_optimizer.step()
else:
beta = self.beta_schedule.get_value(self._n_train_steps_total)
weights = F.softmax(advantage / beta, dim=0)
policy_loss = alpha * log_pi.mean()
if self.use_awr_update and self.weight_loss:
policy_loss = policy_loss + self.awr_weight * (-policy_logpp * len(weights)*weights.detach()).mean()
elif self.use_awr_update:
policy_loss = policy_loss + self.awr_weight * (-policy_logpp).mean()
if self.use_reparam_update:
policy_loss = policy_loss + self.reparam_weight * (-q_new_actions).mean()
policy_loss = self.rl_weight * policy_loss
if self.compute_bc:
train_policy_loss, train_logp_loss, train_mse_loss, _ = self.run_bc_batch(self.demo_train_buffer, self.policy)
policy_loss = policy_loss + self.bc_weight * train_policy_loss
if self.train_bc_on_rl_buffer:
buffer_policy_loss, buffer_train_logp_loss, buffer_train_mse_loss, _ = self.run_bc_batch(self.replay_buffer, self.buffer_policy)
"""
Update networks
"""
if self._n_train_steps_total % self.q_update_period == 0:
self.qf1_optimizer.zero_grad()
qf1_loss.backward()
self.qf1_optimizer.step()
self.qf2_optimizer.zero_grad()
qf2_loss.backward()
self.qf2_optimizer.step()
if self._n_train_steps_total % self.policy_update_period == 0 and self.update_policy:
self.policy_optimizer.zero_grad()
policy_loss.backward()
self.policy_optimizer.step()
if self.train_bc_on_rl_buffer and self._n_train_steps_total % self.policy_update_period == 0 :
self.buffer_policy_optimizer.zero_grad()
buffer_policy_loss.backward()
self.buffer_policy_optimizer.step()
"""
Soft Updates
"""
if self._n_train_steps_total % self.target_update_period == 0:
ptu.soft_update_from_to(
self.qf1, self.target_qf1, self.soft_target_tau
)
ptu.soft_update_from_to(
self.qf2, self.target_qf2, self.soft_target_tau
)
"""
Save some statistics for eval
"""
if self._need_to_update_eval_statistics:
self._need_to_update_eval_statistics = False
"""
Eval should set this to None.
This way, these statistics are only computed for one batch.
"""
policy_loss = (log_pi - q_new_actions).mean()
self.eval_statistics['QF1 Loss'] = np.mean(ptu.get_numpy(qf1_loss))
self.eval_statistics['QF2 Loss'] = np.mean(ptu.get_numpy(qf2_loss))
self.eval_statistics['Policy Loss'] = np.mean(ptu.get_numpy(
policy_loss
))
self.eval_statistics.update(create_stats_ordered_dict(
'Q1 Predictions',
ptu.get_numpy(q1_pred),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Q2 Predictions',
ptu.get_numpy(q2_pred),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Q Targets',
ptu.get_numpy(q_target),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Log Pis',
ptu.get_numpy(log_pi),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Policy mu',
ptu.get_numpy(policy_mean),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Policy log std',
ptu.get_numpy(policy_log_std),
))
self.eval_statistics.update(create_stats_ordered_dict(
'Advantage Weights',
ptu.get_numpy(weights),
))
if self.use_automatic_entropy_tuning:
self.eval_statistics['Alpha'] = alpha.item()
self.eval_statistics['Alpha Loss'] = alpha_loss.item()
if self.compute_bc:
test_policy_loss, test_logp_loss, test_mse_loss, _ = self.run_bc_batch(self.demo_test_buffer, self.policy)
self.eval_statistics.update({
"bc/Train Logprob Loss": ptu.get_numpy(train_logp_loss),
"bc/Test Logprob Loss": ptu.get_numpy(test_logp_loss),
"bc/Train MSE": ptu.get_numpy(train_mse_loss),
"bc/Test MSE": ptu.get_numpy(test_mse_loss),
"bc/train_policy_loss": ptu.get_numpy(train_policy_loss),
"bc/test_policy_loss": ptu.get_numpy(test_policy_loss),
})
if self.train_bc_on_rl_buffer:
test_policy_loss, test_logp_loss, test_mse_loss, _ = self.run_bc_batch(self.replay_buffer,
self.buffer_policy)
_, _, _, _, _, _, _, _, buffer_dist = self.buffer_policy(
obs, reparameterize=True, return_log_prob=True,
)
kldiv = torch.distributions.kl.kl_divergence(dist, buffer_dist)
self.eval_statistics.update({
"buffer_policy/Train Logprob Loss": ptu.get_numpy(buffer_train_logp_loss),
"buffer_policy/Test Logprob Loss": ptu.get_numpy(test_logp_loss),
"buffer_policy/Train MSE": ptu.get_numpy(buffer_train_mse_loss),
"buffer_policy/Test MSE": ptu.get_numpy(test_mse_loss),
"buffer_policy/train_policy_loss": ptu.get_numpy(buffer_policy_loss),
"buffer_policy/test_policy_loss": ptu.get_numpy(test_policy_loss),
"buffer_policy/kl_div":ptu.get_numpy(kldiv.mean()),
})
if self.use_automatic_beta_tuning:
self.eval_statistics.update({
"adaptive_beta/beta":ptu.get_numpy(beta.mean()),
"adaptive_beta/beta loss": ptu.get_numpy(beta_loss.mean()),
})
self._n_train_steps_total += 1
def 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,
)
class ConvVAETrainer(Serializable):
def __init__(
self,
train_dataset,
test_dataset,
model,
batch_size=64,
beta=0.5,
beta_schedule=None,
lr=None,
linearity_weight=0.0,
use_linear_dynamics=False,
noisy_linear_dynamics=False,
scale_linear_dynamics=False,
use_parallel_dataloading=True,
train_data_workers=2,
):
self.quick_init(locals())
self.batch_size = batch_size
self.beta = beta
if lr is None:
lr = 1e-3
self.beta_schedule = beta_schedule
if self.beta_schedule is None:
self.beta_schedule = ConstantSchedule(self.beta)
self.imsize = model.imsize
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.lr = lr
self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr)
self.train_dataset, self.test_dataset = train_dataset, test_dataset
assert self.train_dataset['next_obs'].dtype == np.uint8
assert self.test_dataset['next_obs'].dtype == np.uint8
assert self.train_dataset['obs'].dtype == np.uint8
assert self.test_dataset['obs'].dtype == np.uint8
self.use_parallel_dataloading = use_parallel_dataloading
self.train_data_workers = train_data_workers
self.gaussian_decoder_loss = self.model.gaussian_decoder
if use_parallel_dataloading:
self.train_dataset_pt = ImageDataset(train_dataset,
should_normalize=True)
self.test_dataset_pt = ImageDataset(test_dataset,
should_normalize=True)
self._train_weights = None
base_sampler = InfiniteRandomSampler(self.train_dataset)
self.train_dataloader = DataLoader(
self.train_dataset_pt,
sampler=BatchSampler(
base_sampler,
batch_size=batch_size,
drop_last=False,
),
num_workers=train_data_workers,
pin_memory=True,
)
self.test_dataloader = DataLoader(
self.test_dataset_pt,
sampler=BatchSampler(
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.linearity_weight = linearity_weight
self.use_linear_dynamics = use_linear_dynamics
self.noisy_linear_dynamics = noisy_linear_dynamics
self.scale_linear_dynamics = scale_linear_dynamics
self.vae_logger_stats_for_rl = {}
self._extra_stats_to_log = None
def _kl_np_to_np(self, np_imgs):
torch_input = ptu.np_to_var(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 set_vae(self, vae):
self.model = vae
self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr)
def get_batch(self, train=True):
if self.use_parallel_dataloading:
if not train:
dataloader = self.test_dataloader
else:
dataloader = self.train_dataloader
samples = next(dataloader)
return {
'obs': ptu.Variable(samples[0][0]),
'actions': ptu.Variable(samples[1][0]),
'next_obs': ptu.Variable(samples[2][0]),
}
dataset = self.train_dataset if train else self.test_dataset
ind = np.random.randint(0, len(dataset), self.batch_size)
samples = normalize_image(dataset[ind, :])
return ptu.np_to_var(samples)
def logprob_iwae(self, recon_x, x):
if self.gaussian_decoder_loss:
error = -(recon_x - x)**2
else:
error = x * torch.log(torch.clamp(recon_x, min=1e-30)) \
+ (1-x) * torch.log(torch.clamp(1-recon_x, min=1e-30))
return error
def logprob_vae(self, recon_x, x):
batch_size = recon_x.shape[0]
if self.gaussian_decoder_loss:
return -((recon_x - x)**2).sum() / batch_size
else:
# 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,
) / batch_size
def kl_divergence(self, mu, logvar):
return -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(),
dim=1).mean()
def compute_vae_loss(self, x_recon, x, z_mu, z_logvar, z_sampled, beta):
batch_size = x_recon.shape[0]
k = x_recon.shape[1]
x_recon = x_recon.view((batch_size * k, -1))
x = x.view((batch_size * k, -1))
z_mu = z_mu.view((batch_size * k, -1))
z_logvar = z_logvar.view((batch_size * k, -1))
de = -self.logprob_vae(x_recon, x)
kle = self.kl_divergence(z_mu, z_logvar)
loss = de + beta * kle
return loss, de, kle
def compute_iwae_loss(self, x_recon, x, z_mu, z_logvar, z_sampled, beta):
batch_size = x_recon.shape[0]
log_p_xgz = self.logprob_iwae(x_recon, x).sum(dim=-1)
prior_dist = torch.distributions.Normal(
ptu.Variable(torch.zeros(z_sampled.shape)),
ptu.Variable(torch.ones(z_sampled.shape)))
log_p_z = prior_dist.log_prob(z_sampled).sum(dim=-1)
z_std = torch.exp(0.5 * z_logvar)
encoder_dist = torch.distributions.Normal(z_mu, z_std)
log_q_zgx = encoder_dist.log_prob(z_sampled).sum(dim=-1)
log_w = log_p_xgz + beta * (log_p_z - log_q_zgx)
w_tilde = F.softmax(log_w, dim=-1).detach()
loss = -(log_w * w_tilde).sum() / batch_size
return loss
def state_linearity_loss(self, obs, next_obs, actions):
latent_obs_mu, latent_obs_logvar = self.model.encode(obs)
latent_next_obs_mu, latent_next_obs_logvar = self.model.encode(
next_obs)
if self.noisy_linear_dynamics:
latent_obs = self.model.reparameterize(latent_obs_mu,
latent_obs_logvar)
else:
latent_obs = latent_obs_mu
action_obs_pair = torch.cat([latent_obs, actions], dim=1)
prediction = self.model.linear_constraint_fc(action_obs_pair)
if self.scale_linear_dynamics:
std = latent_next_obs_logvar.mul(0.5).exp_()
scaling = 1 / std
else:
scaling = 1.0
return torch.norm(
scaling * (prediction - latent_next_obs_mu))**2 / self.batch_size
def train_epoch(self, epoch, batches=100):
self.model.own_train()
vae_losses = []
losses = []
des = []
kles = []
linear_losses = []
beta = float(self.beta_schedule.get_value(epoch))
for batch_idx in range(batches):
data = self.get_batch()
obs = data['obs']
next_obs = data['next_obs']
actions = data['actions']
self.optimizer.zero_grad()
x_recon, z_mu, z_logvar, z = self.model(next_obs)
batch_size = x_recon.shape[0]
k = x_recon.shape[1]
x = next_obs.view((batch_size, 1, -1)).repeat(torch.Size([1, k,
1]))
vae_loss, de, kle = self.compute_vae_loss(x_recon, x, z_mu,
z_logvar, z, beta)
loss = vae_loss
'''print("---------------------")
print("loss ", loss.data)
print("z_mu ", z_mu.data[0][0])
print("z_logvar ", z_logvar.data[0][0])
print("z ", z.data[0][0])
print("---------------------")'''
if self.use_linear_dynamics:
linear_dynamics_loss = self.state_linearity_loss(
obs, next_obs, actions)
loss += self.linearity_weight * linear_dynamics_loss
linear_losses.append(float(linear_dynamics_loss.data[0]))
loss.backward()
vae_losses.append(float(vae_loss.data[0]))
losses.append(float(loss.data[0]))
des.append(float(de.data[0]))
kles.append(float(kle.data[0]))
self.optimizer.step()
del data, obs, next_obs, actions, x_recon, z_mu, z_logvar, \
z, x, vae_loss, de, kle, loss
logger.record_tabular("train/epoch", epoch)
logger.record_tabular("train/decoder_loss", np.mean(des))
logger.record_tabular("train/KL", np.mean(kles))
if self.use_linear_dynamics:
logger.record_tabular("train/linear_loss", np.mean(linear_losses))
logger.record_tabular("train/vae_loss", np.mean(vae_losses))
logger.record_tabular("train/loss", np.mean(losses))
def test_epoch(self,
epoch,
save_reconstruction=True,
save_interpolation=True,
save_vae=True):
self.model.eval()
vae_losses = []
iwae_losses = []
losses = []
des = []
kles = []
linear_losses = []
zs = []
beta = float(self.beta_schedule.get_value(epoch))
for batch_idx in range(10):
data = self.get_batch(train=False)
obs = data['obs']
obs = obs.detach()
next_obs = data['next_obs']
next_obs = next_obs.detach()
actions = data['actions']
actions = actions.detach()
x_recon, z_mu, z_logvar, z = self.model(next_obs, n_imp=25)
x_recon = x_recon.detach()
z_mu = z_mu.detach()
z = z.detach()
batch_size = x_recon.shape[0]
k = x_recon.shape[1]
x = next_obs.view((batch_size, 1, -1)).repeat(torch.Size([1, k,
1]))
x = x.detach()
vae_loss, de, kle = self.compute_vae_loss(x_recon, x, z_mu,
z_logvar, z, beta)
vae_loss, de, kle = vae_loss.detach(), de.detach(), kle.detach()
iwae_loss = self.compute_iwae_loss(x_recon, x, z_mu, z_logvar, z,
beta)
iwae_loss = iwae_loss.detach()
loss = vae_loss
if self.use_linear_dynamics:
linear_dynamics_loss = self.state_linearity_loss(
obs, next_obs, actions)
linear_dynamics_loss = linear_dynamics_loss.detach()
loss += self.linearity_weight * linear_dynamics_loss
linear_losses.append(float(
linear_dynamics_loss.data[0])) #here too
z_data = ptu.get_numpy(z_mu[:, 0].cpu())
for i in range(len(z_data)):
zs.append(z_data[i, :].copy())
vae_losses.append(float(vae_loss.data[0]))
iwae_losses.append(float(iwae_loss.data[0]))
losses.append(float(loss.data[0]))
des.append(float(de.data[0]))
kles.append(float(kle.data[0]))
if batch_idx == 0 and save_reconstruction:
n = min(data['next_obs'].size(0), 16)
comparison = torch.cat([
data['next_obs'][:n].narrow(start=0,
length=self.imlength,
dimension=1).contiguous().view(
-1, self.input_channels,
self.imsize, self.imsize),
x_recon[:, 0].contiguous().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)
del comparison
if batch_idx == 0 and save_interpolation:
n = min(data['next_obs'].size(0), 10)
z1 = z_mu[:n, 0]
z2 = z_mu[n:2 * n, 0]
num_steps = 8
z_interp = []
for i in np.linspace(0.0, 1.0, num_steps):
z_interp.append(float(i) * z1 + float(1 - i) * z2)
z_interp = torch.cat(z_interp)
imgs = self.model.decode(z_interp)
imgs = imgs.view((num_steps, n, 3, self.imsize, self.imsize))
imgs = imgs.permute([1, 0, 2, 3, 4])
imgs = imgs.contiguous().view(
(n * num_steps, 3, self.imsize, self.imsize))
save_dir = osp.join(logger.get_snapshot_dir(),
'i_%d.png' % epoch)
save_image(
imgs.data.cpu(),
save_dir,
nrow=num_steps,
)
del imgs
del z_interp
del obs, next_obs, actions, x_recon, z_mu, z_logvar, \
z, x, vae_loss, de, kle, loss
zs = np.array(zs)
self.model.dist_mu = zs.mean(axis=0)
self.model.dist_std = zs.std(axis=0)
del zs
logger.record_tabular("test/decoder_loss", np.mean(des))
logger.record_tabular("test/KL", np.mean(kles))
if self.use_linear_dynamics:
logger.record_tabular("test/linear_loss", np.mean(linear_losses))
logger.record_tabular("test/loss", np.mean(losses))
logger.record_tabular("test/vae_loss", np.mean(vae_losses))
logger.record_tabular("test/iwae_loss", np.mean(iwae_losses))
logger.record_tabular(
"test/iwae_vae_diff",
np.mean(np.array(iwae_losses) - np.array(vae_losses)))
logger.record_tabular("beta", beta)
process = psutil.Process(os.getpid())
logger.record_tabular("RAM Usage (Mb)",
int(process.memory_info().rss / 1000000))
num_active_dims = 0
num_active_dims2 = 0
for std in self.model.dist_std:
if std > 0.15:
num_active_dims += 1
if std > 0.05:
num_active_dims2 += 1
logger.record_tabular("num_active_dims", num_active_dims)
logger.record_tabular("num_active_dims2", num_active_dims2)
logger.dump_tabular()
if save_vae:
logger.save_itr_params(epoch,
self.model,
prefix='vae',
save_anyway=True) # slow...