def experiment(variant):
rdim = variant["rdim"]
vae_paths = {
2:
"/home/ashvin/data/s3doodad/ashvin/vae/point2d-conv-sweep2/run1/id0/params.pkl",
4:
"/home/ashvin/data/s3doodad/ashvin/vae/point2d-conv-sweep2/run1/id1/params.pkl",
8:
"/home/ashvin/data/s3doodad/ashvin/vae/point2d-conv-sweep2/run1/id2/params.pkl",
16:
"/home/ashvin/data/s3doodad/ashvin/vae/point2d-conv-sweep2/run1/id3/params.pkl"
}
vae_path = vae_paths[rdim]
vae = torch.load(vae_path)
print("loaded", vae_path)
if variant['multitask']:
env = MultitaskImagePoint2DEnv(**variant['env_kwargs'])
env = VAEWrappedImageGoalEnv(
env,
vae,
use_vae_obs=True,
use_vae_reward=True,
use_vae_goals=True,
)
# render_goals=True, render_rollouts=True)
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"]
use_env_goals = variant["use_env_goals"]
vae_path = variant["vae_paths"][str(rdim)]
render = variant["render"]
wrap_mujoco_env = variant.get("wrap_mujoco_env", False)
# vae = torch.load(vae_path)
# print("loaded", vae_path)
from railrl.envs.wrappers import ImageMujocoEnv, NormalizedBoxEnv
from railrl.images.camera import sawyer_init_camera
env = variant["env"](**variant['env_kwargs'])
env = NormalizedBoxEnv(ImageMujocoEnv(
env,
imsize=84,
keep_prev=0,
init_camera=sawyer_init_camera,
))
if wrap_mujoco_env:
env = ImageMujocoEnv(env, 84, camera_name="topview", transpose=True, normalize=True)
if use_env_goals:
track_qpos_goal = variant.get("track_qpos_goal", 0)
env = VAEWrappedImageGoalEnv(env, vae_path, use_vae_obs=True,
use_vae_reward=True, use_vae_goals=True,
render_goals=render, render_rollouts=render, track_qpos_goal=track_qpos_goal)
else:
env = VAEWrappedEnv(env, vae_path, use_vae_obs=True,
use_vae_reward=True, use_vae_goals=True,
render_goals=render, render_rollouts=render)
env = MultitaskToFlatEnv(env)
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']
)
algorithm.to(ptu.device)
env._wrapped_env.vae.to(ptu.device)