def experiment(variant):
env = variant['env_class'](**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
goal_dim = env.goal_dim
qf1 = ConcatMlp(
input_size=obs_dim + action_dim + goal_dim,
output_size=1,
hidden_sizes=[400, 300],
)
qf2 = ConcatMlp(
input_size=obs_dim + action_dim + goal_dim,
output_size=1,
hidden_sizes=[400, 300],
)
policy = TanhMlpPolicy(
input_size=obs_dim + goal_dim,
output_size=action_dim,
hidden_sizes=[400, 300],
)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
replay_buffer = variant['replay_buffer_class'](
env=env, **variant['replay_buffer_kwargs'])
algorithm = HerTd3(env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
replay_buffer=replay_buffer,
**variant['algo_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
env = CylinderXYPusher2DEnv(**variant['env_kwargs'])
if variant['normalize']:
env = NormalizedBoxEnv(env)
es = EpsilonGreedy(
action_space=env.action_space,
prob_random_action=0.1,
)
obs_dim = env.observation_space.low.size
action_dim = env.action_space.low.size
goal_dim = env.goal_dim
qf1 = ConcatMlp(
input_size=obs_dim + action_dim + goal_dim,
output_size=1,
hidden_sizes=[400, 300],
)
qf2 = ConcatMlp(
input_size=obs_dim + action_dim + goal_dim,
output_size=1,
hidden_sizes=[400, 300],
)
policy = TanhMlpPolicy(
input_size=obs_dim + goal_dim,
output_size=action_dim,
hidden_sizes=[400, 300],
)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
replay_buffer = SimpleHerReplayBuffer(env=env,
**variant['replay_buffer_kwargs'])
algorithm = HerTd3(env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
replay_buffer=replay_buffer,
**variant['algo_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def HER_baseline_td3_experiment(variant):
import rlkit.torch.pytorch_util as ptu
from rlkit.data_management.obs_dict_replay_buffer import \
ObsDictRelabelingBuffer
from rlkit.exploration_strategies.base import (
PolicyWrappedWithExplorationStrategy)
from rlkit.torch.her.her_td3 import HerTd3
from rlkit.torch.networks import MergedCNN, CNNPolicy
import torch
from multiworld.core.image_env import ImageEnv
from rlkit.misc.asset_loader import load_local_or_remote_file
init_camera = variant.get("init_camera", None)
presample_goals = variant.get('presample_goals', False)
presampled_goals_path = get_presampled_goals_path(
variant.get('presampled_goals_path', None))
if 'env_id' in variant:
import gym
import multiworld
multiworld.register_all_envs()
env = gym.make(variant['env_id'])
else:
env = variant["env_class"](**variant['env_kwargs'])
image_env = ImageEnv(
env,
variant.get('imsize'),
reward_type='image_sparse',
init_camera=init_camera,
transpose=True,
normalize=True,
)
if presample_goals:
if presampled_goals_path is None:
image_env.non_presampled_goal_img_is_garbage = True
presampled_goals = variant['generate_goal_dataset_fctn'](
env=image_env, **variant['goal_generation_kwargs'])
else:
presampled_goals = load_local_or_remote_file(
presampled_goals_path).item()
del image_env
env = ImageEnv(
env,
variant.get('imsize'),
reward_type='image_distance',
init_camera=init_camera,
transpose=True,
normalize=True,
presampled_goals=presampled_goals,
)
else:
env = image_env
es = get_exploration_strategy(variant, env)
observation_key = variant.get('observation_key', 'image_observation')
desired_goal_key = variant.get('desired_goal_key', 'image_desired_goal')
achieved_goal_key = desired_goal_key.replace("desired", "achieved")
imsize = variant['imsize']
action_dim = env.action_space.low.size
qf1 = MergedCNN(input_width=imsize,
input_height=imsize,
output_size=1,
input_channels=3 * 2,
added_fc_input_size=action_dim,
**variant['cnn_params'])
qf2 = MergedCNN(input_width=imsize,
input_height=imsize,
output_size=1,
input_channels=3 * 2,
added_fc_input_size=action_dim,
**variant['cnn_params'])
policy = CNNPolicy(
input_width=imsize,
input_height=imsize,
added_fc_input_size=0,
output_size=action_dim,
input_channels=3 * 2,
output_activation=torch.tanh,
**variant['cnn_params'],
)
target_qf1 = MergedCNN(input_width=imsize,
input_height=imsize,
output_size=1,
input_channels=3 * 2,
added_fc_input_size=action_dim,
**variant['cnn_params'])
target_qf2 = MergedCNN(input_width=imsize,
input_height=imsize,
output_size=1,
input_channels=3 * 2,
added_fc_input_size=action_dim,
**variant['cnn_params'])
target_policy = CNNPolicy(
input_width=imsize,
input_height=imsize,
added_fc_input_size=0,
output_size=action_dim,
input_channels=3 * 2,
output_activation=torch.tanh,
**variant['cnn_params'],
)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
replay_buffer = ObsDictRelabelingBuffer(
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']
algo_kwargs['replay_buffer'] = replay_buffer
base_kwargs = algo_kwargs['base_kwargs']
base_kwargs['training_env'] = env
base_kwargs['render'] = variant["render"]
base_kwargs['render_during_eval'] = variant["render"]
her_kwargs = algo_kwargs['her_kwargs']
her_kwargs['observation_key'] = observation_key
her_kwargs['desired_goal_key'] = desired_goal_key
algorithm = HerTd3(env,
qf1=qf1,
qf2=qf2,
policy=policy,
target_qf1=target_qf1,
target_qf2=target_qf2,
target_policy=target_policy,
exploration_policy=exploration_policy,
**variant['algo_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def her_td3_experiment(variant):
env = variant['env_class'](**variant['env_kwargs'])
if 'history_len' in variant:
history_len = variant['history_len']
env = MultiTaskHistoryEnv(env, history_len=history_len)
if variant.get('make_silent_env', True):
env = MultitaskEnvToSilentMultitaskEnv(env)
if variant['normalize']:
env = NormalizedBoxEnv(env)
exploration_type = variant['exploration_type']
if exploration_type == 'ou':
es = OUStrategy(
action_space=env.action_space,
**variant['es_kwargs']
)
elif exploration_type == 'gaussian':
es = GaussianStrategy(
action_space=env.action_space,
**variant['es_kwargs'],
)
elif exploration_type == 'epsilon':
es = EpsilonGreedy(
action_space=env.action_space,
**variant['es_kwargs'],
)
else:
raise Exception("Invalid type: " + exploration_type)
obs_dim = env.observation_space.low.size
action_dim = env.action_space.low.size
goal_dim = env.goal_space.low.size
qf1 = ConcatMlp(
input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs']
)
qf2 = ConcatMlp(
input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs']
)
policy = TanhMlpPolicy(
input_size=obs_dim + goal_dim,
output_size=action_dim,
**variant['policy_kwargs']
)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
replay_buffer = variant['replay_buffer_class'](
env=env,
**variant['replay_buffer_kwargs']
)
algorithm = HerTd3(
env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
replay_buffer=replay_buffer,
**variant['algo_kwargs']
)
algorithm.to(ptu.device)
algorithm.train()
def her_td3_experiment(variant):
env = variant['env_class'](**variant['env_kwargs'])
observation_key = variant.get('observation_key', 'observation')
desired_goal_key = variant.get('desired_goal_key', 'desired_goal')
replay_buffer = ObsDictRelabelingBuffer(env=env,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
**variant['replay_buffer_kwargs'])
obs_dim = env.observation_space.spaces['observation'].low.size
action_dim = env.action_space.low.size
goal_dim = env.observation_space.spaces['desired_goal'].low.size
if variant['normalize']:
env = NormalizedBoxEnv(env)
exploration_type = variant['exploration_type']
if exploration_type == 'ou':
es = OUStrategy(action_space=env.action_space,
max_sigma=0.1,
**variant['es_kwargs'])
elif exploration_type == 'gaussian':
es = GaussianStrategy(
action_space=env.action_space,
max_sigma=0.1,
min_sigma=0.1, # Constant sigma
**variant['es_kwargs'],
)
elif exploration_type == 'epsilon':
es = EpsilonGreedy(
action_space=env.action_space,
prob_random_action=0.1,
**variant['es_kwargs'],
)
else:
raise Exception("Invalid type: " + exploration_type)
qf1 = ConcatMlp(input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs'])
qf2 = ConcatMlp(input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs'])
policy = TanhMlpPolicy(input_size=obs_dim + goal_dim,
output_size=action_dim,
**variant['policy_kwargs'])
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
algorithm = HerTd3(env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
replay_buffer=replay_buffer,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
**variant['algo_kwargs'])
if ptu.gpu_enabled():
qf1.to(ptu.device)
qf2.to(ptu.device)
policy.to(ptu.device)
algorithm.to(ptu.device)
algorithm.train()
def her_td3_experiment(variant):
import gym
import multiworld.envs.mujoco
import multiworld.envs.pygame
import rlkit.samplers.rollout_functions as rf
import rlkit.torch.pytorch_util as ptu
from rlkit.exploration_strategies.base import (
PolicyWrappedWithExplorationStrategy)
from rlkit.exploration_strategies.epsilon_greedy import EpsilonGreedy
from rlkit.exploration_strategies.gaussian_strategy import GaussianStrategy
from rlkit.exploration_strategies.ou_strategy import OUStrategy
from rlkit.torch.grill.launcher import get_video_save_func
from rlkit.torch.her.her_td3 import HerTd3
from rlkit.torch.networks import ConcatMlp, TanhMlpPolicy
from rlkit.data_management.obs_dict_replay_buffer import (
ObsDictRelabelingBuffer)
if 'env_id' in variant:
env = gym.make(variant['env_id'])
else:
env = variant['env_class'](**variant['env_kwargs'])
observation_key = variant['observation_key']
desired_goal_key = variant['desired_goal_key']
variant['algo_kwargs']['her_kwargs']['observation_key'] = observation_key
variant['algo_kwargs']['her_kwargs']['desired_goal_key'] = desired_goal_key
if variant.get('normalize', False):
raise NotImplementedError()
achieved_goal_key = desired_goal_key.replace("desired", "achieved")
replay_buffer = ObsDictRelabelingBuffer(
env=env,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
achieved_goal_key=achieved_goal_key,
**variant['replay_buffer_kwargs'])
obs_dim = env.observation_space.spaces['observation'].low.size
action_dim = env.action_space.low.size
goal_dim = env.observation_space.spaces['desired_goal'].low.size
exploration_type = variant['exploration_type']
if exploration_type == 'ou':
es = OUStrategy(action_space=env.action_space, **variant['es_kwargs'])
elif exploration_type == 'gaussian':
es = GaussianStrategy(
action_space=env.action_space,
**variant['es_kwargs'],
)
elif exploration_type == 'epsilon':
es = EpsilonGreedy(
action_space=env.action_space,
**variant['es_kwargs'],
)
else:
raise Exception("Invalid type: " + exploration_type)
qf1 = ConcatMlp(input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs'])
qf2 = ConcatMlp(input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs'])
policy = TanhMlpPolicy(input_size=obs_dim + goal_dim,
output_size=action_dim,
**variant['policy_kwargs'])
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
algorithm = HerTd3(env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
replay_buffer=replay_buffer,
**variant['algo_kwargs'])
if variant.get("save_video", 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,
policy,
variant,
)
algorithm.post_epoch_funcs.append(video_func)
algorithm.to(ptu.device)
algorithm.train()
def grill_her_td3_experiment(variant):
env = variant["env_class"](**variant['env_kwargs'])
render = variant["render"]
rdim = variant["rdim"]
vae_path = variant["vae_paths"][str(rdim)]
reward_params = variant.get("reward_params", dict())
init_camera = variant.get("init_camera", None)
if init_camera is None:
camera_name = "topview"
else:
camera_name = None
env = ImageEnv(
env,
84,
init_camera=init_camera,
camera_name=camera_name,
transpose=True,
normalize=True,
)
env = VAEWrappedEnv(env,
vae_path,
decode_goals=render,
render_goals=render,
render_rollouts=render,
reward_params=reward_params,
**variant.get('vae_wrapped_env_kwargs', {}))
if variant['normalize']:
env = NormalizedBoxEnv(env)
exploration_type = variant['exploration_type']
exploration_noise = variant.get('exploration_noise', 0.1)
if exploration_type == 'ou':
es = OUStrategy(action_space=env.action_space)
elif exploration_type == 'gaussian':
es = GaussianStrategy(
action_space=env.action_space,
max_sigma=exploration_noise,
min_sigma=exploration_noise, # Constant sigma
)
elif exploration_type == 'epsilon':
es = EpsilonGreedy(
action_space=env.action_space,
prob_random_action=exploration_noise,
)
else:
raise Exception("Invalid type: " + exploration_type)
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
hidden_sizes = variant.get('hidden_sizes', [400, 300])
qf1 = ConcatMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=hidden_sizes,
)
qf2 = ConcatMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=hidden_sizes,
)
policy = TanhMlpPolicy(
input_size=obs_dim,
output_size=action_dim,
hidden_sizes=hidden_sizes,
)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
training_mode = variant.get("training_mode", "train")
testing_mode = variant.get("testing_mode", "test")
testing_env = pickle.loads(pickle.dumps(env))
testing_env.mode(testing_mode)
training_env = pickle.loads(pickle.dumps(env))
training_env.mode(training_mode)
relabeling_env = pickle.loads(pickle.dumps(env))
relabeling_env.mode(training_mode)
relabeling_env.disable_render()
video_vae_env = pickle.loads(pickle.dumps(env))
video_vae_env.mode("video_vae")
video_goal_env = pickle.loads(pickle.dumps(env))
video_goal_env.mode("video_env")
replay_buffer = ObsDictRelabelingBuffer(
env=relabeling_env,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
achieved_goal_key=achieved_goal_key,
**variant['replay_kwargs'])
variant["algo_kwargs"]["replay_buffer"] = replay_buffer
algorithm = HerTd3(testing_env,
training_env=training_env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
render=render,
render_during_eval=render,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
**variant['algo_kwargs'])
if ptu.gpu_enabled():
print("using GPU")
algorithm.to(ptu.device)
for e in [testing_env, training_env, video_vae_env, video_goal_env]:
e.vae.to(ptu.device)
algorithm.train()
if variant.get("save_video", True):
logdir = logger.get_snapshot_dir()
policy.train(False)
filename = osp.join(logdir, 'video_final_env.mp4')
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,
)
dump_video(video_goal_env, policy, filename, rollout_function)
filename = osp.join(logdir, 'video_final_vae.mp4')
dump_video(video_vae_env, policy, filename, rollout_function)