def experiment(variant):
from railrl.core import logger
import railrl.torch.pytorch_util as ptu
beta = variant["beta"]
representation_size = variant["representation_size"]
train_data, test_data, info = generate_vae_dataset(
**variant['get_data_kwargs'])
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
m = ConvVAE(representation_size, input_channels=3)
if ptu.gpu_enabled():
m.to(ptu.device)
gpu_id = variant.get("gpu_id", None)
if gpu_id is not None:
ptu.set_device(gpu_id)
t = ConvVAETrainer(train_data,
test_data,
m,
beta=beta,
beta_schedule=beta_schedule,
**variant['algo_kwargs'])
save_period = variant['save_period']
for epoch in range(variant['num_epochs']):
should_save_imgs = (epoch % save_period == 0)
t.train_epoch(epoch)
t.test_epoch(epoch,
save_reconstruction=should_save_imgs,
save_scatterplot=should_save_imgs)
if should_save_imgs:
t.dump_samples(epoch)
def experiment(variant):
beta = variant["beta"]
representation_size = variant["representation_size"]
train_data, test_data = get_data(10000)
m = ConvVAE(representation_size)
t = ConvVAETrainer(train_data, test_data, m, beta=beta, use_cuda=False)
for epoch in range(10):
t.train_epoch(epoch)
t.test_epoch(epoch)
t.dump_samples(epoch)
def experiment(variant):
if variant["use_gpu"]:
ptu.set_gpu_mode(True)
beta = variant["beta"]
representation_size = variant["representation_size"]
m = ConvVAE(representation_size, input_channels=3)
t = ConvVAETrainer(train_data, test_data, m, beta=beta)
for epoch in range(1001):
t.train_epoch(epoch)
t.test_epoch(epoch)
t.dump_samples(epoch)
def experiment(variant):
from railrl.core import logger
import railrl.torch.pytorch_util as ptu
beta = variant["beta"]
representation_size = variant["representation_size"]
# train_data, test_data, info = generate_vae_dataset(
# **variant['get_data_kwargs']
# )
num_divisions = 5
images = np.zeros((num_divisions * 10000, 21168))
for i in range(num_divisions):
imgs = np.load(
'/home/murtaza/vae_data/sawyer_torque_control_images100000_' +
str(i + 1) + '.npy')
images[i * 10000:(i + 1) * 10000] = imgs
print(i)
mid = int(num_divisions * 10000 * .9)
train_data, test_data = images[:mid], images[mid:]
info = dict()
logger.save_extra_data(info)
logger.get_snapshot_dir()
if 'beta_schedule_kwargs' in variant:
kwargs = variant['beta_schedule_kwargs']
kwargs['y_values'][2] = variant['beta']
kwargs['x_values'][1] = variant['flat_x']
kwargs['x_values'][2] = variant['ramp_x'] + variant['flat_x']
beta_schedule = PiecewiseLinearSchedule(
**variant['beta_schedule_kwargs'])
else:
beta_schedule = None
m = ConvVAE(representation_size,
input_channels=3,
**variant['conv_vae_kwargs'])
if ptu.gpu_enabled():
m.cuda()
t = ConvVAETrainer(train_data,
test_data,
m,
beta=beta,
beta_schedule=beta_schedule,
**variant['algo_kwargs'])
save_period = variant['save_period']
for epoch in range(variant['num_epochs']):
should_save_imgs = (epoch % save_period == 0)
t.train_epoch(epoch)
t.test_epoch(epoch,
save_reconstruction=should_save_imgs,
save_scatterplot=should_save_imgs)
if should_save_imgs:
t.dump_samples(epoch)
def experiment(variant):
if variant["use_gpu"]:
gpu_id = variant["gpu_id"]
ptu.set_gpu_mode(True)
ptu.set_device(gpu_id)
beta = variant["beta"]
representation_size = variant["representation_size"]
train_data, test_data = get_data(10000)
m = ConvVAE(representation_size, input_channels=3)
t = ConvVAETrainer(train_data, test_data, m, beta=beta, use_cuda=True)
for epoch in range(50):
t.train_epoch(epoch)
t.test_epoch(epoch)
t.dump_samples(epoch)
def experiment(variant):
if variant["use_gpu"]:
gpu_id = variant["gpu_id"]
ptu.set_gpu_mode(True)
ptu.set_device(gpu_id)
beta = variant["beta"]
representation_size = variant["representation_size"]
train_data, test_data = get_data(10000)
m = ConvVAE(representation_size)
t = ConvVAETrainer(train_data, test_data, m, beta=beta, do_scatterplot=False)
for epoch in range(101):
t.train_epoch(epoch)
t.test_epoch(epoch)
t.dump_samples(epoch)
def experiment(variant):
num_feat_points=variant['feat_points']
from railrl.core import logger
beta = variant["beta"]
print('collecting data')
train_data, test_data, info = get_data(**variant['get_data_kwargs'])
print('finish collecting data')
logger.save_extra_data(info)
logger.get_snapshot_dir()
m = SpatialAutoEncoder(2 * num_feat_points, num_feat_points, input_channels=3)
# m = ConvVAE(2*num_feat_points, input_channels=3)
t = ConvVAETrainer(train_data, test_data, m, lr=variant['lr'], beta=beta)
for epoch in range(variant['num_epochs']):
t.train_epoch(epoch)
t.test_epoch(epoch)
t.dump_samples(epoch)
def experiment(variant):
from railrl.core import logger
import railrl.torch.pytorch_util as ptu
beta = variant["beta"]
representation_size = variant["representation_size"]
#this has both states and images so can't use generate vae dataset
X = np.load(
'/home/murtaza/vae_data/sawyer_torque_control_ou_imgs_zoomed_out10000.npy'
)
Y = np.load(
'/home/murtaza/vae_data/sawyer_torque_control_ou_states_zoomed_out10000.npy'
)
Y = np.concatenate((Y[:, :7], Y[:, 14:]), axis=1)
X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=.1)
info = dict()
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
m = ConvVAE(representation_size,
input_channels=3,
state_sim_debug=True,
state_size=Y.shape[1],
**variant['conv_vae_kwargs'])
if ptu.gpu_enabled():
m.cuda()
t = ConvVAETrainer((X_train, Y_train), (X_test, Y_test),
m,
beta=beta,
beta_schedule=beta_schedule,
state_sim_debug=True,
**variant['algo_kwargs'])
save_period = variant['save_period']
for epoch in range(variant['num_epochs']):
should_save_imgs = (epoch % save_period == 0)
t.train_epoch(epoch)
t.test_epoch(epoch,
save_reconstruction=should_save_imgs,
save_scatterplot=should_save_imgs)
if should_save_imgs:
t.dump_samples(epoch)
def experiment(variant):
if variant["use_gpu"]:
gpu_id = variant["gpu_id"]
ptu.set_gpu_mode(True)
ptu.set_device(gpu_id)
beta = variant["beta"]
representation_size = variant["representation_size"]
train_data, test_data = get_data(10000)
m = ConvVAE(representation_size, input_channels=3)
t = ConvVAETrainer(train_data,
test_data,
m,
beta_schedule=PiecewiseLinearSchedule([0, 400, 800],
[0.5, 0.5, beta]))
for epoch in range(1001):
t.train_epoch(epoch)
t.test_epoch(epoch)
t.dump_samples(epoch)
def experiment(variant):
from railrl.core import logger
import railrl.torch.pytorch_util as ptu
beta = variant["beta"]
representation_size = variant["representation_size"]
train_data, test_data, info = generate_vae_dataset(
**variant['generate_vae_dataset_kwargs'])
logger.save_extra_data(info)
logger.get_snapshot_dir()
if 'beta_schedule_kwargs' in variant:
# kwargs = variant['beta_schedule_kwargs']
# kwargs['y_values'][2] = variant['beta']
# kwargs['x_values'][1] = variant['flat_x']
# kwargs['x_values'][2] = variant['ramp_x'] + variant['flat_x']
variant['beta_schedule_kwargs']['y_values'][-1] = variant['beta']
beta_schedule = PiecewiseLinearSchedule(
**variant['beta_schedule_kwargs'])
else:
beta_schedule = None
m = ConvVAE(representation_size,
input_channels=3,
**variant['conv_vae_kwargs'])
if ptu.gpu_enabled():
m.cuda()
t = ConvVAETrainer(train_data,
test_data,
m,
beta=beta,
beta_schedule=beta_schedule,
**variant['algo_kwargs'])
save_period = variant['save_period']
for epoch in range(variant['num_epochs']):
should_save_imgs = (epoch % save_period == 0)
t.train_epoch(epoch)
t.test_epoch(epoch,
save_reconstruction=should_save_imgs,
save_scatterplot=should_save_imgs)
if should_save_imgs:
t.dump_samples(epoch)
def experiment(variant):
from railrl.core import logger
import railrl.torch.pytorch_util as ptu
beta = variant["beta"]
representation_size = variant["representation_size"]
train_data, test_data, info = get_data(**variant['get_data_kwargs'])
logger.save_extra_data(info)
logger.get_snapshot_dir()
beta_schedule = PiecewiseLinearSchedule(**variant['beta_schedule_kwargs'])
m = ConvVAE(representation_size, input_channels=3)
if ptu.gpu_enabled():
m.to(ptu.device)
t = ConvVAETrainer(train_data,
test_data,
m,
beta=beta,
beta_schedule=beta_schedule,
**variant['algo_kwargs'])
for epoch in range(variant['num_epochs']):
t.train_epoch(epoch)
t.test_epoch(epoch)
t.dump_samples(epoch)
def experiment(variant):
from railrl.core import logger
import railrl.torch.pytorch_util as ptu
beta = variant["beta"]
representation_size = variant["representation_size"]
train_data, test_data, info = variant['generate_vae_dataset_fn'](
variant['generate_vae_dataset_kwargs']
)
logger.save_extra_data(info)
logger.get_snapshot_dir()
if 'beta_schedule_kwargs' in variant:
# kwargs = variant['beta_schedule_kwargs']
# kwargs['y_values'][2] = variant['beta']
# kwargs['x_values'][1] = variant['flat_x']
# kwargs['x_values'][2] = variant['ramp_x'] + variant['flat_x']
variant['beta_schedule_kwargs']['y_values'][-1] = variant['beta']
beta_schedule = PiecewiseLinearSchedule(**variant['beta_schedule_kwargs'])
else:
beta_schedule = None
m = variant['vae'](representation_size, **variant['vae_kwargs'])
m.to(ptu.device)
t = ConvVAETrainer(train_data, test_data, m, beta=beta,
beta_schedule=beta_schedule, **variant['algo_kwargs'])
save_period = variant['save_period']
for epoch in range(variant['num_epochs']):
should_save_imgs = (epoch % save_period == 0)
t.train_epoch(epoch)
t.test_epoch(epoch, save_reconstruction=should_save_imgs,
save_scatterplot=should_save_imgs)
if should_save_imgs:
t.dump_samples(epoch)
if variant['dump_skew_debug_plots']:
t.dump_best_reconstruction(epoch)
t.dump_worst_reconstruction(epoch)
t.dump_sampling_histogram(epoch)
t.update_train_weights()
def td3_experiment_online_vae_exploring(variant):
import railrl.samplers.rollout_functions as rf
import railrl.torch.pytorch_util as ptu
from railrl.data_management.online_vae_replay_buffer import \
OnlineVaeRelabelingBuffer
from railrl.exploration_strategies.base import (
PolicyWrappedWithExplorationStrategy)
from railrl.torch.her.online_vae_joint_algo import OnlineVaeHerJointAlgo
from railrl.torch.networks import FlattenMlp, TanhMlpPolicy
from railrl.torch.td3.td3 import TD3
from railrl.torch.vae.vae_trainer import ConvVAETrainer
preprocess_rl_variant(variant)
env = get_envs(variant)
es = get_exploration_strategy(variant, env)
observation_key = variant.get('observation_key', 'latent_observation')
desired_goal_key = variant.get('desired_goal_key', 'latent_desired_goal')
achieved_goal_key = desired_goal_key.replace("desired", "achieved")
obs_dim = (env.observation_space.spaces[observation_key].low.size +
env.observation_space.spaces[desired_goal_key].low.size)
action_dim = env.action_space.low.size
qf1 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs'],
)
qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs'],
)
policy = TanhMlpPolicy(
input_size=obs_dim,
output_size=action_dim,
**variant['policy_kwargs'],
)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
exploring_qf1 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs'],
)
exploring_qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs'],
)
exploring_policy = TanhMlpPolicy(
input_size=obs_dim,
output_size=action_dim,
**variant['policy_kwargs'],
)
exploring_exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=exploring_policy,
)
vae = env.vae
replay_buffer = OnlineVaeRelabelingBuffer(
vae=vae,
env=env,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
achieved_goal_key=achieved_goal_key,
**variant['replay_buffer_kwargs'])
variant["algo_kwargs"]["replay_buffer"] = replay_buffer
if variant.get('use_replay_buffer_goals', False):
env.replay_buffer = replay_buffer
env.use_replay_buffer_goals = True
vae_trainer_kwargs = variant.get('vae_trainer_kwargs')
t = ConvVAETrainer(variant['vae_train_data'],
variant['vae_test_data'],
vae,
beta=variant['online_vae_beta'],
**vae_trainer_kwargs)
control_algorithm = TD3(env=env,
training_env=env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
**variant['algo_kwargs'])
exploring_algorithm = TD3(env=env,
training_env=env,
qf1=exploring_qf1,
qf2=exploring_qf2,
policy=exploring_policy,
exploration_policy=exploring_exploration_policy,
**variant['algo_kwargs'])
assert 'vae_training_schedule' not in variant,\
"Just put it in joint_algo_kwargs"
algorithm = OnlineVaeHerJointAlgo(vae=vae,
vae_trainer=t,
env=env,
training_env=env,
policy=policy,
exploration_policy=exploration_policy,
replay_buffer=replay_buffer,
algo1=control_algorithm,
algo2=exploring_algorithm,
algo1_prefix="Control_",
algo2_prefix="VAE_Exploration_",
observation_key=observation_key,
desired_goal_key=desired_goal_key,
**variant['joint_algo_kwargs'])
algorithm.to(ptu.device)
vae.to(ptu.device)
if variant.get("save_video", True):
policy.train(False)
rollout_function = rf.create_rollout_function(
rf.multitask_rollout,
max_path_length=algorithm.max_path_length,
observation_key=algorithm.observation_key,
desired_goal_key=algorithm.desired_goal_key,
)
video_func = get_video_save_func(
rollout_function,
env,
algorithm.eval_policy,
variant,
)
algorithm.post_train_funcs.append(video_func)
algorithm.train()
def train_vae(variant, return_data=False):
from railrl.misc.ml_util import PiecewiseLinearSchedule
from railrl.torch.vae.conv_vae import (
ConvVAE,
SpatialAutoEncoder,
AutoEncoder,
)
import railrl.torch.vae.conv_vae as conv_vae
from railrl.torch.vae.vae_trainer import ConvVAETrainer
from railrl.core import logger
import railrl.torch.pytorch_util as ptu
from railrl.pythonplusplus import identity
import torch
beta = variant["beta"]
representation_size = variant["representation_size"]
train_data, test_data, info = generate_vae_dataset_from_demos(
variant['generate_vae_dataset_kwargs'])
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 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):
m = AutoEncoder(representation_size,
decoder_output_activation=decoder_activation,
**variant['vae_kwargs'])
elif variant.get('use_spatial_auto_encoder', False):
raise NotImplementedError(
'This is currently broken, please update SpatialAutoEncoder then remove this line'
)
m = SpatialAutoEncoder(representation_size,
int(representation_size / 2))
else:
vae_class = variant.get('vae_class', ConvVAE)
m = vae_class(representation_size,
decoder_output_activation=decoder_activation,
**variant['vae_kwargs'])
m.to(ptu.device)
t = ConvVAETrainer(train_data,
test_data,
m,
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)
t.train_epoch(epoch)
t.test_epoch(
epoch,
save_reconstruction=should_save_imgs,
save_scatterplot=should_save_imgs,
# save_vae=False,
)
if should_save_imgs:
t.dump_samples(epoch)
if dump_skew_debug_plots:
t.dump_best_reconstruction(epoch)
t.dump_worst_reconstruction(epoch)
t.dump_sampling_histogram(epoch)
t.update_train_weights()
logger.save_extra_data(m, 'vae.pkl', mode='pickle')
if return_data:
return m, train_data, test_data
return m
def tdm_td3_experiment_online_vae(variant):
import railrl.samplers.rollout_functions as rf
import railrl.torch.pytorch_util as ptu
from railrl.data_management.online_vae_replay_buffer import \
OnlineVaeRelabelingBuffer
from railrl.exploration_strategies.base import (
PolicyWrappedWithExplorationStrategy)
from railrl.state_distance.tdm_networks import TdmQf, TdmPolicy
from railrl.torch.vae.vae_trainer import ConvVAETrainer
from railrl.torch.online_vae.online_vae_tdm_td3 import OnlineVaeTdmTd3
preprocess_rl_variant(variant)
env = get_envs(variant)
es = get_exploration_strategy(variant, env)
vae_trainer_kwargs = variant.get('vae_trainer_kwargs')
observation_key = variant.get('observation_key', 'latent_observation')
desired_goal_key = variant.get('desired_goal_key', 'latent_desired_goal')
achieved_goal_key = desired_goal_key.replace("desired", "achieved")
obs_dim = (env.observation_space.spaces[observation_key].low.size)
goal_dim = (env.observation_space.spaces[desired_goal_key].low.size)
action_dim = env.action_space.low.size
vectorized = 'vectorized' in env.reward_type
variant['algo_kwargs']['tdm_td3_kwargs']['tdm_kwargs'][
'vectorized'] = vectorized
norm_order = env.norm_order
# variant['algo_kwargs']['tdm_td3_kwargs']['tdm_kwargs'][
# 'norm_order'] = norm_order
qf1 = TdmQf(env=env,
vectorized=vectorized,
norm_order=norm_order,
observation_dim=obs_dim,
goal_dim=goal_dim,
action_dim=action_dim,
**variant['qf_kwargs'])
qf2 = TdmQf(env=env,
vectorized=vectorized,
norm_order=norm_order,
observation_dim=obs_dim,
goal_dim=goal_dim,
action_dim=action_dim,
**variant['qf_kwargs'])
policy = TdmPolicy(env=env,
observation_dim=obs_dim,
goal_dim=goal_dim,
action_dim=action_dim,
**variant['policy_kwargs'])
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
vae = env.vae
replay_buffer = OnlineVaeRelabelingBuffer(
vae=vae,
env=env,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
achieved_goal_key=achieved_goal_key,
**variant['replay_buffer_kwargs'])
algo_kwargs = variant['algo_kwargs']['tdm_td3_kwargs']
td3_kwargs = algo_kwargs['td3_kwargs']
td3_kwargs['training_env'] = env
tdm_kwargs = algo_kwargs['tdm_kwargs']
tdm_kwargs['observation_key'] = observation_key
tdm_kwargs['desired_goal_key'] = desired_goal_key
algo_kwargs["replay_buffer"] = replay_buffer
t = ConvVAETrainer(variant['vae_train_data'],
variant['vae_test_data'],
vae,
beta=variant['online_vae_beta'],
**vae_trainer_kwargs)
render = variant["render"]
assert 'vae_training_schedule' not in variant, "Just put it in algo_kwargs"
algorithm = OnlineVaeTdmTd3(
online_vae_kwargs=dict(vae=vae,
vae_trainer=t,
**variant['algo_kwargs']['online_vae_kwargs']),
tdm_td3_kwargs=dict(env=env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
**variant['algo_kwargs']['tdm_td3_kwargs']),
)
algorithm.to(ptu.device)
vae.to(ptu.device)
if variant.get("save_video", True):
policy.train(False)
rollout_function = rf.create_rollout_function(
rf.tdm_rollout,
init_tau=algorithm._sample_max_tau_for_rollout(),
decrement_tau=algorithm.cycle_taus_for_rollout,
cycle_tau=algorithm.cycle_taus_for_rollout,
max_path_length=algorithm.max_path_length,
observation_key=algorithm.observation_key,
desired_goal_key=algorithm.desired_goal_key,
)
video_func = get_video_save_func(
rollout_function,
env,
algorithm.eval_policy,
variant,
)
algorithm.post_train_funcs.append(video_func)
algorithm.to(ptu.device)
if not variant.get("do_state_exp", False):
env.vae.to(ptu.device)
algorithm.train()
def get_n_train_vae(latent_dim,
env,
vae_train_epochs,
num_image_examples,
vae_kwargs,
vae_trainer_kwargs,
vae_architecture,
vae_save_period=10,
vae_test_p=.9,
decoder_activation='sigmoid',
vae_class='VAE',
**kwargs):
env.goal_sampling_mode = 'test'
image_examples = unnormalize_image(
env.sample_goals(num_image_examples)['desired_goal'])
n = int(num_image_examples * vae_test_p)
train_dataset = ImageObservationDataset(image_examples[:n, :])
test_dataset = ImageObservationDataset(image_examples[n:, :])
if decoder_activation == 'sigmoid':
decoder_activation = torch.nn.Sigmoid()
vae_class = vae_class.lower()
if vae_class == 'VAE'.lower():
vae_class = ConvVAE
elif vae_class == 'SpatialVAE'.lower():
vae_class = SpatialAutoEncoder
else:
raise RuntimeError("Invalid VAE Class: {}".format(vae_class))
vae = vae_class(latent_dim,
architecture=vae_architecture,
decoder_output_activation=decoder_activation,
**vae_kwargs)
trainer = ConvVAETrainer(vae, **vae_trainer_kwargs)
logger.remove_tabular_output('progress.csv', relative_to_snapshot_dir=True)
logger.add_tabular_output('vae_progress.csv',
relative_to_snapshot_dir=True)
for epoch in range(vae_train_epochs):
should_save_imgs = (epoch % vae_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)
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, vae)
logger.save_extra_data(vae, 'vae.pkl', mode='pickle')
logger.remove_tabular_output('vae_progress.csv',
relative_to_snapshot_dir=True)
logger.add_tabular_output('progress.csv', relative_to_snapshot_dir=True)
return vae
def experiment(variant):
from railrl.core import logger
import railrl.torch.pytorch_util as ptu
beta = variant["beta"]
representation_size = variant["representation_size"]
train_data, test_data, info = variant['generate_vae_dataset_fn'](
variant['generate_vae_dataset_kwargs']
)
uniform_dataset=generate_uniform_dataset_reacher(
**variant['generate_uniform_dataset_kwargs']
)
logger.save_extra_data(info)
logger.get_snapshot_dir()
beta_schedule = None
m = variant['vae'](representation_size, decoder_output_activation=nn.Sigmoid(), **variant['vae_kwargs'])
m.to(ptu.device)
t = ConvVAETrainer(train_data, test_data, m, beta=beta,
beta_schedule=beta_schedule, **variant['algo_kwargs'])
save_period = variant['save_period']
for epoch in range(variant['num_epochs']):
should_save_imgs = (epoch % save_period == 0)
t.train_epoch(epoch)
t.log_loss_under_uniform(m, uniform_dataset)
t.test_epoch(epoch, save_reconstruction=should_save_imgs,
save_scatterplot=should_save_imgs)
if should_save_imgs:
t.dump_samples(epoch)
if variant['dump_skew_debug_plots']:
t.dump_best_reconstruction(epoch)
t.dump_worst_reconstruction(epoch)
t.dump_sampling_histogram(epoch)
t.dump_uniform_imgs_and_reconstructions(dataset=uniform_dataset, epoch=epoch)
if epoch % variant['train_weight_update_period'] == 0:
t.update_train_weights()
def experiment(variant):
from railrl.core import logger
import railrl.torch.pytorch_util as ptu
beta = variant["beta"]
representation_size = variant["representation_size"]
train_data, test_data, info = variant['generate_vae_dataset_fn'](
variant['generate_vae_dataset_kwargs'])
uniform_dataset = load_local_or_remote_file(
variant['uniform_dataset_path']).item()
uniform_dataset = unormalize_image(uniform_dataset['image_desired_goal'])
logger.save_extra_data(info)
logger.get_snapshot_dir()
if 'beta_schedule_kwargs' in variant:
# kwargs = variant['beta_schedule_kwargs']
# kwargs['y_values'][2] = variant['beta']
# kwargs['x_values'][1] = variant['flat_x']
# kwargs['x_values'][2] = variant['ramp_x'] + variant['flat_x']
variant['beta_schedule_kwargs']['y_values'][-1] = variant['beta']
beta_schedule = PiecewiseLinearSchedule(
**variant['beta_schedule_kwargs'])
else:
beta_schedule = None
m = variant['vae'](representation_size,
decoder_output_activation=nn.Sigmoid(),
**variant['vae_kwargs'])
m.to(ptu.device)
t = ConvVAETrainer(train_data,
test_data,
m,
beta=beta,
beta_schedule=beta_schedule,
**variant['algo_kwargs'])
save_period = variant['save_period']
for epoch in range(variant['num_epochs']):
should_save_imgs = (epoch % save_period == 0)
t.train_epoch(epoch)
t.log_loss_under_uniform(
m, uniform_dataset,
variant['algo_kwargs']['priority_function_kwargs'])
t.test_epoch(epoch,
save_reconstruction=should_save_imgs,
save_scatterplot=should_save_imgs)
if should_save_imgs:
t.dump_samples(epoch)
if variant['dump_skew_debug_plots']:
t.dump_best_reconstruction(epoch)
t.dump_worst_reconstruction(epoch)
t.dump_sampling_histogram(epoch)
t.dump_uniform_imgs_and_reconstructions(
dataset=uniform_dataset, epoch=epoch)
t.update_train_weights()
def experiment(variant):
from railrl.core import logger
import railrl.torch.pytorch_util as ptu
beta = variant["beta"]
representation_size = variant["representation_size"]
data = joblib.load(variant['file'])
obs = data['obs']
size = int(data['size'])
dataset = obs[:size, :]
n = int(size * .9)
train_data = dataset[:n, :]
test_data = dataset[n:, :]
logger.get_snapshot_dir()
print('SIZE: ', size)
uniform_dataset = generate_uniform_dataset_door(
**variant['generate_uniform_dataset_kwargs']
)
logger.get_snapshot_dir()
if 'beta_schedule_kwargs' in variant:
# kwargs = variant['beta_schedule_kwargs']
# kwargs['y_values'][2] = variant['beta']
# kwargs['x_values'][1] = variant['flat_x']
# kwargs['x_values'][2] = variant['ramp_x'] + variant['flat_x']
variant['beta_schedule_kwargs']['y_values'][-1] = variant['beta']
beta_schedule = PiecewiseLinearSchedule(**variant['beta_schedule_kwargs'])
else:
beta_schedule = None
m = variant['vae'](representation_size, decoder_output_activation=nn.Sigmoid(), **variant['vae_kwargs'])
m.to(ptu.device)
t = ConvVAETrainer(train_data, test_data, m, beta=beta,
beta_schedule=beta_schedule, **variant['algo_kwargs'])
save_period = variant['save_period']
for epoch in range(variant['num_epochs']):
should_save_imgs = (epoch % save_period == 0)
t.train_epoch(epoch)
t.log_loss_under_uniform(uniform_dataset)
t.test_epoch(epoch, save_reconstruction=should_save_imgs,
save_scatterplot=should_save_imgs)
if should_save_imgs:
t.dump_samples(epoch)
if variant['dump_skew_debug_plots']:
t.dump_best_reconstruction(epoch)
t.dump_worst_reconstruction(epoch)
t.dump_sampling_histogram(epoch)
t.dump_uniform_imgs_and_reconstructions(dataset=uniform_dataset, epoch=epoch)
t.update_train_weights()
