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)