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):
if variant['multitask']:
env = MultitaskFullVAEPoint2DEnv(
**variant['env_kwargs']) # used point2d-conv-sweep/run1/id4
env = MultitaskToFlatEnv(env)
# else:
# env = Pusher2DEnv(**variant['env_kwargs'])
if variant['normalize']:
env = NormalizedBoxEnv(env)
exploration_type = variant['exploration_type']
if exploration_type == 'ou':
es = OUStrategy(action_space=env.action_space)
elif exploration_type == 'gaussian':
es = GaussianStrategy(
action_space=env.action_space,
max_sigma=0.1,
min_sigma=0.1, # Constant sigma
)
elif exploration_type == 'epsilon':
es = EpsilonGreedy(
action_space=env.action_space,
prob_random_action=0.1,
)
else:
raise Exception("Invalid type: " + exploration_type)
obs_dim = env.observation_space.low.size
action_dim = env.action_space.low.size
qf1 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[400, 300],
)
qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[400, 300],
)
policy = TanhMlpPolicy(
input_size=obs_dim,
output_size=action_dim,
hidden_sizes=[400, 300],
)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
algorithm = TD3(env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
**variant['algo_kwargs'])
if variant["use_gpu"]:
gpu_id = variant["gpu_id"]
ptu.set_gpu_mode(True)
ptu.set_device(gpu_id)
algorithm.to(ptu.device)
env._wrapped_env.vae.to(ptu.device)
algorithm.train()
def experiment(variant):
if variant["use_gpu"]:
ptu.set_gpu_mode(True)
ptu.set_device(0)
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)
for epoch in range(1000):
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, use_cuda=True)
for epoch in range(10):
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(train_data,
test_data,
representation_size,
beta=beta,
use_cuda=True,
input_channels=3)
for epoch in range(50):
m.train_epoch(epoch)
m.test_epoch(epoch)
m.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):
rdim = variant["rdim"]
vae_paths = {
2:
"/home/ashvin/data/s3doodad/ashvin/vae/point2d-conv-sweep2/run0/id1/params.pkl",
4:
"/home/ashvin/data/s3doodad/ashvin/vae/point2d-conv-sweep2/run0/id4/params.pkl"
}
vae_path = vae_paths[rdim]
vae = joblib.load(vae_path)
print("loaded", vae_path)
if variant['multitask']:
env = MultitaskImagePoint2DEnv(**variant['env_kwargs'])
env = VAEWrappedEnv(env,
vae,
use_vae_obs=True,
use_vae_reward=False,
use_vae_goals=False)
env = MultitaskToFlatEnv(env)
# else:
# env = Pusher2DEnv(**variant['env_kwargs'])
if variant['normalize']:
env = NormalizedBoxEnv(env)
exploration_type = variant['exploration_type']
if exploration_type == 'ou':
es = OUStrategy(action_space=env.action_space)
elif exploration_type == 'gaussian':
es = GaussianStrategy(
action_space=env.action_space,
max_sigma=0.1,
min_sigma=0.1, # Constant sigma
)
elif exploration_type == 'epsilon':
es = EpsilonGreedy(
action_space=env.action_space,
prob_random_action=0.1,
)
else:
raise Exception("Invalid type: " + exploration_type)
obs_dim = env.observation_space.low.size
action_dim = env.action_space.low.size
qf1 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[400, 300],
)
qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[400, 300],
)
policy = TanhMlpPolicy(
input_size=obs_dim,
output_size=action_dim,
hidden_sizes=[400, 300],
)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
algorithm = TD3(env,
training_env=env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
**variant['algo_kwargs'])
print("use_gpu", variant["use_gpu"], bool(variant["use_gpu"]))
if variant["use_gpu"]:
gpu_id = variant["gpu_id"]
ptu.set_gpu_mode(True)
ptu.set_device(gpu_id)
algorithm.to(ptu.device)
env._wrapped_env.vae.to(ptu.device)
algorithm.train()
def experiment(variant):
rdim = variant["rdim"]
vae_paths = {
2:
"/home/ashvin/data/s3doodad/ashvin/vae/new-pusher2d/run2/id0/params.pkl",
4:
"/home/ashvin/data/s3doodad/ashvin/vae/new-pusher2d/run2/id1/params.pkl",
8:
"/home/ashvin/data/s3doodad/ashvin/vae/new-pusher2d/run2/id2/params.pkl",
16:
"/home/ashvin/data/s3doodad/ashvin/vae/new-pusher2d/run2/id3/params.pkl"
}
vae_path = vae_paths[rdim]
vae = torch.load(vae_path)
print("loaded", vae_path)
if variant['multitask']:
env = FullPusher2DEnv(**variant["env_kwargs"])
env = ImageMujocoEnv(env,
84,
camera_name="topview",
transpose=True,
normalize=True)
env = VAEWrappedImageGoalEnv(env,
vae,
use_vae_obs=True,
use_vae_reward=True,
use_vae_goals=True,
render_goals=True,
render_rollouts=True,
track_qpos_goal=5)
env = MultitaskToFlatEnv(env)
# else:
# env = Pusher2DEnv(**variant['env_kwargs'])
if variant['normalize']:
env = NormalizedBoxEnv(env)
exploration_type = variant['exploration_type']
if exploration_type == 'ou':
es = OUStrategy(action_space=env.action_space)
elif exploration_type == 'gaussian':
es = GaussianStrategy(
action_space=env.action_space,
max_sigma=0.1,
min_sigma=0.1, # Constant sigma
)
elif exploration_type == 'epsilon':
es = EpsilonGreedy(
action_space=env.action_space,
prob_random_action=0.1,
)
else:
raise Exception("Invalid type: " + exploration_type)
obs_dim = env.observation_space.low.size
action_dim = env.action_space.low.size
qf1 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[400, 300],
)
qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[400, 300],
)
policy = TanhMlpPolicy(
input_size=obs_dim,
output_size=action_dim,
hidden_sizes=[400, 300],
)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
algorithm = TD3(env,
training_env=env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
**variant['algo_kwargs'])
print("use_gpu", variant["use_gpu"], bool(variant["use_gpu"]))
if variant["use_gpu"]:
gpu_id = variant["gpu_id"]
ptu.set_gpu_mode(True)
ptu.set_device(gpu_id)
algorithm.to(ptu.device)
env._wrapped_env.vae.to(ptu.device)
algorithm.train()
