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,
is_auto_encoder=variant['is_auto_encoder'],
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 = 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):
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)