def experiment(variant):
render = variant.get("render", False)
debug = variant.get("debug", False)
if variant.get("pretrained_algorithm_path", False):
resume(variant)
return
env_class = variant["env_class"]
env_kwargs = variant["env_kwargs"]
expl_env = env_class(**env_kwargs)
eval_env = env_class(**env_kwargs)
env = eval_env
if variant.get('sparse_reward', False):
expl_env = RewardWrapperEnv(expl_env, compute_hand_sparse_reward)
eval_env = RewardWrapperEnv(eval_env, compute_hand_sparse_reward)
if variant.get('add_env_demos', False):
variant["path_loader_kwargs"]["demo_paths"].append(variant["env_demo_path"])
if variant.get('add_env_offpolicy_data', False):
variant["path_loader_kwargs"]["demo_paths"].append(variant["env_offpolicy_data_path"])
if variant.get("use_masks", False):
mask_wrapper_kwargs = variant.get("mask_wrapper_kwargs", dict())
expl_mask_distribution_kwargs = variant["expl_mask_distribution_kwargs"]
expl_mask_distribution = DiscreteDistribution(**expl_mask_distribution_kwargs)
expl_env = RewardMaskWrapper(env, expl_mask_distribution, **mask_wrapper_kwargs)
eval_mask_distribution_kwargs = variant["eval_mask_distribution_kwargs"]
eval_mask_distribution = DiscreteDistribution(**eval_mask_distribution_kwargs)
eval_env = RewardMaskWrapper(env, eval_mask_distribution, **mask_wrapper_kwargs)
env = eval_env
path_loader_kwargs = variant.get("path_loader_kwargs", {})
stack_obs = path_loader_kwargs.get("stack_obs", 1)
if stack_obs > 1:
expl_env = StackObservationEnv(expl_env, stack_obs=stack_obs)
eval_env = StackObservationEnv(eval_env, stack_obs=stack_obs)
observation_key = variant.get('observation_key', 'latent_observation')
desired_goal_key = variant.get('desired_goal_key', 'latent_desired_goal')
achieved_goal_key = variant.get('achieved_goal_key', 'latent_achieved_goal')
obs_dim = (
env.observation_space.spaces[observation_key].low.size
+ env.observation_space.spaces[desired_goal_key].low.size
)
action_dim = eval_env.action_space.low.size
if hasattr(expl_env, 'info_sizes'):
env_info_sizes = expl_env.info_sizes
else:
env_info_sizes = dict()
replay_buffer_kwargs=dict(
env=env,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
achieved_goal_key=achieved_goal_key,
)
replay_buffer_kwargs.update(variant.get('replay_buffer_kwargs', dict()))
replay_buffer = ConcatToObsWrapper(
ObsDictRelabelingBuffer(**replay_buffer_kwargs),
["resampled_goals", ],
)
replay_buffer_kwargs.update(variant.get('demo_replay_buffer_kwargs', dict()))
demo_train_buffer = ConcatToObsWrapper(
ObsDictRelabelingBuffer(**replay_buffer_kwargs),
["resampled_goals", ],
)
demo_test_buffer = ConcatToObsWrapper(
ObsDictRelabelingBuffer(**replay_buffer_kwargs),
["resampled_goals", ],
)
M = variant['layer_size']
qf1 = ConcatMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[M, M],
)
qf2 = ConcatMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[M, M],
)
target_qf1 = ConcatMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[M, M],
)
target_qf2 = ConcatMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[M, M],
)
policy_class = variant.get("policy_class", TanhGaussianPolicy)
policy_kwargs = variant['policy_kwargs']
policy_path = variant.get("policy_path", False)
if policy_path:
policy = load_local_or_remote_file(policy_path)
else:
policy = policy_class(
obs_dim=obs_dim,
action_dim=action_dim,
**policy_kwargs,
)
buffer_policy_path = variant.get("buffer_policy_path", False)
if buffer_policy_path:
buffer_policy = load_local_or_remote_file(buffer_policy_path)
else:
buffer_policy_class = variant.get("buffer_policy_class", policy_class)
buffer_policy = buffer_policy_class(
obs_dim=obs_dim,
action_dim=action_dim,
**variant.get("buffer_policy_kwargs", policy_kwargs),
)
expl_policy = policy
exploration_kwargs = variant.get('exploration_kwargs', {})
if exploration_kwargs:
if exploration_kwargs.get("deterministic_exploration", False):
expl_policy = MakeDeterministic(policy)
exploration_strategy = exploration_kwargs.get("strategy", None)
if exploration_strategy is None:
pass
elif exploration_strategy == 'ou':
es = OUStrategy(
action_space=expl_env.action_space,
max_sigma=exploration_kwargs['noise'],
min_sigma=exploration_kwargs['noise'],
)
expl_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=expl_policy,
)
elif exploration_strategy == 'gauss_eps':
es = GaussianAndEpislonStrategy(
action_space=expl_env.action_space,
max_sigma=exploration_kwargs['noise'],
min_sigma=exploration_kwargs['noise'], # constant sigma
epsilon=0,
)
expl_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=expl_policy,
)
else:
error
trainer = AWACTrainer(
env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
buffer_policy=buffer_policy,
**variant['trainer_kwargs']
)
if variant['collection_mode'] == 'online':
expl_path_collector = MdpStepCollector(
expl_env,
policy,
)
algorithm = TorchOnlineRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
max_path_length=variant['max_path_length'],
batch_size=variant['batch_size'],
num_epochs=variant['num_epochs'],
num_eval_steps_per_epoch=variant['num_eval_steps_per_epoch'],
num_expl_steps_per_train_loop=variant['num_expl_steps_per_train_loop'],
num_trains_per_train_loop=variant['num_trains_per_train_loop'],
min_num_steps_before_training=variant['min_num_steps_before_training'],
)
else:
eval_path_collector = GoalConditionedPathCollector(
eval_env,
MakeDeterministic(policy),
observation_key=observation_key,
desired_goal_key=desired_goal_key,
render=render,
)
expl_path_collector = GoalConditionedPathCollector(
expl_env,
expl_policy,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
render=render,
)
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
max_path_length=variant['max_path_length'],
batch_size=variant['batch_size'],
num_epochs=variant['num_epochs'],
num_eval_steps_per_epoch=variant['num_eval_steps_per_epoch'],
num_expl_steps_per_train_loop=variant['num_expl_steps_per_train_loop'],
num_trains_per_train_loop=variant['num_trains_per_train_loop'],
min_num_steps_before_training=variant['min_num_steps_before_training'],
)
algorithm.to(ptu.device)
if variant.get("save_video", False):
renderer_kwargs = variant.get("renderer_kwargs", {})
save_video_kwargs = variant.get("save_video_kwargs", {})
def get_video_func(
env,
policy,
tag,
):
renderer = EnvRenderer(**renderer_kwargs)
state_goal_distribution = GoalDictDistributionFromMultitaskEnv(
env,
desired_goal_keys=[desired_goal_key],
)
image_goal_distribution = AddImageDistribution(
env=env,
base_distribution=state_goal_distribution,
image_goal_key='image_desired_goal',
renderer=renderer,
)
img_env = InsertImageEnv(env, renderer=renderer)
rollout_function = partial(
rf.multitask_rollout,
max_path_length=variant['max_path_length'],
observation_key=observation_key,
desired_goal_key=desired_goal_key,
return_dict_obs=True,
)
reward_fn = ContextualRewardFnFromMultitaskEnv(
env=env,
achieved_goal_from_observation=IndexIntoAchievedGoal(observation_key),
desired_goal_key=desired_goal_key,
achieved_goal_key="state_achieved_goal",
)
contextual_env = ContextualEnv(
img_env,
context_distribution=image_goal_distribution,
reward_fn=reward_fn,
observation_key=observation_key,
)
video_func = get_save_video_function(
rollout_function,
contextual_env,
policy,
tag=tag,
imsize=renderer.width,
image_format='CWH',
**save_video_kwargs
)
return video_func
expl_video_func = get_video_func(expl_env, expl_policy, "expl")
eval_video_func = get_video_func(eval_env, MakeDeterministic(policy), "eval")
algorithm.post_train_funcs.append(eval_video_func)
algorithm.post_train_funcs.append(expl_video_func)
if variant.get('save_paths', False):
algorithm.post_train_funcs.append(save_paths)
if variant.get('load_demos', False):
path_loader_class = variant.get('path_loader_class', MDPPathLoader)
path_loader = path_loader_class(trainer,
replay_buffer=replay_buffer,
demo_train_buffer=demo_train_buffer,
demo_test_buffer=demo_test_buffer,
**path_loader_kwargs
)
path_loader.load_demos()
if variant.get('pretrain_policy', False):
trainer.pretrain_policy_with_bc(
policy,
demo_train_buffer,
demo_test_buffer,
trainer.bc_num_pretrain_steps,
)
if variant.get('pretrain_rl', False):
trainer.pretrain_q_with_bc_data()
if variant.get('save_pretrained_algorithm', False):
p_path = osp.join(logger.get_snapshot_dir(), 'pretrain_algorithm.p')
pt_path = osp.join(logger.get_snapshot_dir(), 'pretrain_algorithm.pt')
data = algorithm._get_snapshot()
data['algorithm'] = algorithm
torch.save(data, open(pt_path, "wb"))
torch.save(data, open(p_path, "wb"))
algorithm.train()
def td3_experiment(variant):
import rlkit.samplers.rollout_functions as rf
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.td3.td3 import TD3 as TD3Trainer
from rlkit.torch.torch_rl_algorithm import TorchBatchRLAlgorithm
from rlkit.torch.networks import ConcatMlp, TanhMlpPolicy
# preprocess_rl_variant(variant)
env = get_envs(variant)
expl_env = env
eval_env = env
es = get_exploration_strategy(variant, env)
if variant.get("use_masks", False):
mask_wrapper_kwargs = variant.get("mask_wrapper_kwargs", dict())
expl_mask_distribution_kwargs = variant[
"expl_mask_distribution_kwargs"]
expl_mask_distribution = DiscreteDistribution(
**expl_mask_distribution_kwargs)
expl_env = RewardMaskWrapper(env, expl_mask_distribution,
**mask_wrapper_kwargs)
eval_mask_distribution_kwargs = variant[
"eval_mask_distribution_kwargs"]
eval_mask_distribution = DiscreteDistribution(
**eval_mask_distribution_kwargs)
eval_env = RewardMaskWrapper(env, eval_mask_distribution,
**mask_wrapper_kwargs)
env = eval_env
max_path_length = variant['max_path_length']
observation_key = variant.get('observation_key', 'latent_observation')
desired_goal_key = variant.get('desired_goal_key', 'latent_desired_goal')
achieved_goal_key = variant.get('achieved_goal_key',
'latent_achieved_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 = ConcatMlp(input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
qf2 = ConcatMlp(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'])
target_qf1 = ConcatMlp(input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
target_qf2 = ConcatMlp(input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
target_policy = TanhMlpPolicy(input_size=obs_dim,
output_size=action_dim,
**variant['policy_kwargs'])
if variant.get("use_subgoal_policy", False):
from rlkit.policies.timed_policy import SubgoalPolicyWrapper
subgoal_policy_kwargs = variant.get('subgoal_policy_kwargs', {})
policy = SubgoalPolicyWrapper(wrapped_policy=policy,
env=env,
episode_length=max_path_length,
**subgoal_policy_kwargs)
target_policy = SubgoalPolicyWrapper(wrapped_policy=target_policy,
env=env,
episode_length=max_path_length,
**subgoal_policy_kwargs)
expl_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,
# use_masks=variant.get("use_masks", False),
**variant['replay_buffer_kwargs'])
trainer = TD3Trainer(policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
target_policy=target_policy,
**variant['td3_trainer_kwargs'])
# if variant.get("use_masks", False):
# from rlkit.torch.her.her import MaskedHERTrainer
# trainer = MaskedHERTrainer(trainer)
# else:
trainer = HERTrainer(trainer)
if variant.get("do_state_exp", False):
eval_path_collector = GoalConditionedPathCollector(
eval_env,
policy,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
# use_masks=variant.get("use_masks", False),
# full_mask=True,
)
expl_path_collector = GoalConditionedPathCollector(
expl_env,
expl_policy,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
# use_masks=variant.get("use_masks", False),
)
else:
eval_path_collector = VAEWrappedEnvPathCollector(
env,
policy,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
goal_sampling_mode=['evaluation_goal_sampling_mode'],
)
expl_path_collector = VAEWrappedEnvPathCollector(
env,
expl_policy,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
goal_sampling_mode=['exploration_goal_sampling_mode'],
)
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=env,
evaluation_env=env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
max_path_length=max_path_length,
**variant['algo_kwargs'])
vis_variant = variant.get('vis_kwargs', {})
vis_list = vis_variant.get('vis_list', [])
if variant.get("save_video", True):
if variant.get("do_state_exp", False):
rollout_function = rf.create_rollout_function(
rf.multitask_rollout,
max_path_length=max_path_length,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
# use_masks=variant.get("use_masks", False),
# full_mask=True,
# vis_list=vis_list,
)
video_func = get_video_save_func(
rollout_function,
env,
policy,
variant,
)
else:
video_func = VideoSaveFunction(
env,
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 twin_sac_experiment(variant):
import rlkit.torch.pytorch_util as ptu
from rlkit.data_management.obs_dict_replay_buffer import \
ObsDictRelabelingBuffer
from rlkit.torch.networks import ConcatMlp
from rlkit.torch.sac.policies import TanhGaussianPolicy
from rlkit.torch.torch_rl_algorithm import TorchBatchRLAlgorithm
from rlkit.torch.sac.policies import MakeDeterministic
from rlkit.torch.sac.sac import SACTrainer
preprocess_rl_variant(variant)
env = get_envs(variant)
max_path_length = variant['max_path_length']
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 = ConcatMlp(input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
qf2 = ConcatMlp(input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
target_qf1 = ConcatMlp(input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
target_qf2 = ConcatMlp(input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
policy = TanhGaussianPolicy(obs_dim=obs_dim,
action_dim=action_dim,
**variant['policy_kwargs'])
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'])
trainer = SACTrainer(env=env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**variant['twin_sac_trainer_kwargs'])
trainer = HERTrainer(trainer)
if variant.get("do_state_exp", False):
eval_path_collector = GoalConditionedPathCollector(
env,
MakeDeterministic(policy),
observation_key=observation_key,
desired_goal_key=desired_goal_key,
)
expl_path_collector = GoalConditionedPathCollector(
env,
policy,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
)
else:
eval_path_collector = VAEWrappedEnvPathCollector(
variant['evaluation_goal_sampling_mode'],
env,
MakeDeterministic(policy),
observation_key=observation_key,
desired_goal_key=desired_goal_key,
)
expl_path_collector = VAEWrappedEnvPathCollector(
variant['exploration_goal_sampling_mode'],
env,
policy,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
)
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=env,
evaluation_env=env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
max_path_length=max_path_length,
**variant['algo_kwargs'])
if variant.get("save_video", True):
video_func = VideoSaveFunction(
env,
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 relabeling_tsac_experiment(variant):
if 'presample_goals' in variant:
raise NotImplementedError()
if 'env_id' in variant:
eval_env = gym.make(variant['env_id'])
expl_env = gym.make(variant['env_id'])
else:
eval_env = variant['env_class'](**variant['env_kwargs'])
expl_env = variant['env_class'](**variant['env_kwargs'])
observation_key = variant['observation_key']
desired_goal_key = variant['desired_goal_key']
if variant.get('normalize', False):
raise NotImplementedError()
achieved_goal_key = desired_goal_key.replace("desired", "achieved")
replay_buffer = ObsDictRelabelingBuffer(
env=eval_env,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
achieved_goal_key=achieved_goal_key,
**variant['replay_buffer_kwargs'])
obs_dim = eval_env.observation_space.spaces['observation'].low.size
action_dim = eval_env.action_space.low.size
goal_dim = eval_env.observation_space.spaces['desired_goal'].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'])
target_qf1 = ConcatMlp(input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs'])
target_qf2 = ConcatMlp(input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs'])
policy = TanhGaussianPolicy(obs_dim=obs_dim + goal_dim,
action_dim=action_dim,
**variant['policy_kwargs'])
max_path_length = variant['max_path_length']
eval_policy = MakeDeterministic(policy)
trainer = SACTrainer(env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**variant['twin_sac_trainer_kwargs'])
trainer = HERTrainer(trainer)
eval_path_collector = GoalConditionedPathCollector(
eval_env,
eval_policy,
max_path_length,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
)
expl_path_collector = GoalConditionedPathCollector(
expl_env,
policy,
max_path_length,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
)
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
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()