def tdm_td3_experiment(variant):
env = variant['env_class'](**variant['env_kwargs'])
tdm_normalizer = None
qf1 = TdmQf(
env=env,
vectorized=True,
tdm_normalizer=tdm_normalizer,
**variant['qf_kwargs']
)
qf2 = TdmQf(
env=env,
vectorized=True,
tdm_normalizer=tdm_normalizer,
**variant['qf_kwargs']
)
policy = TdmPolicy(
env=env,
tdm_normalizer=tdm_normalizer,
**variant['policy_kwargs']
)
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)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
replay_buffer = variant['replay_buffer_class'](
env=env,
**variant['replay_buffer_kwargs']
)
qf_criterion = variant['qf_criterion_class']()
algo_kwargs = variant['algo_kwargs']
algo_kwargs['td3_kwargs']['qf_criterion'] = qf_criterion
algo_kwargs['tdm_kwargs']['tdm_normalizer'] = tdm_normalizer
algorithm = TdmTd3(
env,
qf1=qf1,
qf2=qf2,
replay_buffer=replay_buffer,
policy=policy,
exploration_policy=exploration_policy,
**algo_kwargs
)
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
env = variant['env_class'](**variant['env_kwargs'])
# env = NormalizedBoxEnv(env)
# tdm_normalizer = TdmNormalizer(
# env,
# vectorized=True,
# max_tau=variant['algo_kwargs']['tdm_kwargs']['max_tau'],
# )
tdm_normalizer = None
qf = TdmQf(env=env,
vectorized=True,
tdm_normalizer=tdm_normalizer,
**variant['qf_kwargs'])
policy = TdmPolicy(env=env,
tdm_normalizer=tdm_normalizer,
**variant['policy_kwargs'])
es = OUStrategy(action_space=env.action_space, **variant['es_kwargs'])
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
replay_buffer = HerReplayBuffer(env=env,
**variant['her_replay_buffer_kwargs'])
qf_criterion = variant['qf_criterion_class']()
ddpg_tdm_kwargs = variant['algo_kwargs']
ddpg_tdm_kwargs['ddpg_kwargs']['qf_criterion'] = qf_criterion
ddpg_tdm_kwargs['tdm_kwargs']['tdm_normalizer'] = tdm_normalizer
algorithm = TdmDdpg(env,
qf=qf,
replay_buffer=replay_buffer,
policy=policy,
exploration_policy=exploration_policy,
**variant['algo_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
env = variant['env_class'](**variant['env_kwargs'])
tdm_normalizer = None
qf1 = TdmQf(
env=env,
vectorized=True,
tdm_normalizer=tdm_normalizer,
**variant['qf_kwargs']
)
qf2 = TdmQf(
env=env,
vectorized=True,
tdm_normalizer=tdm_normalizer,
**variant['qf_kwargs']
)
policy = TdmPolicy(
env=env,
tdm_normalizer=tdm_normalizer,
**variant['policy_kwargs']
)
es = OUStrategy(
action_space=env.action_space,
**variant['es_kwargs']
)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
replay_buffer = HerReplayBuffer(
env=env,
**variant['her_replay_buffer_kwargs']
)
qf_criterion = variant['qf_criterion_class']()
algo_kwargs = variant['algo_kwargs']
algo_kwargs['td3_kwargs']['qf_criterion'] = qf_criterion
algo_kwargs['tdm_kwargs']['tdm_normalizer'] = tdm_normalizer
algorithm = TdmTd3(
env,
qf1=qf1,
qf2=qf2,
replay_buffer=replay_buffer,
policy=policy,
exploration_policy=exploration_policy,
**variant['algo_kwargs']
)
algorithm.to(ptu.device)
algorithm.train()
def tdm_twin_sac_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.state_distance.tdm_networks import (
TdmQf,
TdmVf,
StochasticTdmPolicy,
)
from rlkit.state_distance.tdm_twin_sac import TdmTwinSAC
preprocess_rl_variant(variant)
env = get_envs(variant)
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)
goal_dim = (env.observation_space.spaces[desired_goal_key].low.size)
action_dim = env.action_space.low.size
vectorized = 'vectorized' in env.reward_type
norm_order = env.norm_order
variant['algo_kwargs']['tdm_kwargs']['vectorized'] = vectorized
variant['qf_kwargs']['vectorized'] = vectorized
variant['vf_kwargs']['vectorized'] = vectorized
variant['qf_kwargs']['norm_order'] = norm_order
variant['vf_kwargs']['norm_order'] = norm_order
qf1 = TdmQf(env=env,
observation_dim=obs_dim,
goal_dim=goal_dim,
action_dim=action_dim,
**variant['qf_kwargs'])
qf2 = TdmQf(env=env,
observation_dim=obs_dim,
goal_dim=goal_dim,
action_dim=action_dim,
**variant['qf_kwargs'])
vf = TdmVf(env=env,
observation_dim=obs_dim,
goal_dim=goal_dim,
**variant['vf_kwargs'])
policy = StochasticTdmPolicy(env=env,
observation_dim=obs_dim,
goal_dim=goal_dim,
action_dim=action_dim,
**variant['policy_kwargs'])
variant['replay_buffer_kwargs']['vectorized'] = vectorized
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"]
tdm_kwargs = algo_kwargs['tdm_kwargs']
tdm_kwargs['observation_key'] = observation_key
tdm_kwargs['desired_goal_key'] = desired_goal_key
algorithm = TdmTwinSAC(env,
qf1=qf1,
qf2=qf2,
vf=vf,
policy=policy,
**variant['algo_kwargs'])
if variant.get("save_video", True):
rollout_function = rf.create_rollout_function(
rf.tdm_rollout,
init_tau=algorithm._sample_max_tau_for_rollout(),
decrement_tau=algorithm.cycle_taus_for_rollout,
cycle_tau=algorithm.cycle_taus_for_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,
algorithm.eval_policy,
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 tdm_td3_experiment_online_vae(variant):
import rlkit.samplers.rollout_functions as rf
import rlkit.torch.pytorch_util as ptu
from rlkit.data_management.online_vae_replay_buffer import \
OnlineVaeRelabelingBuffer
from rlkit.exploration_strategies.base import (
PolicyWrappedWithExplorationStrategy)
from rlkit.state_distance.tdm_networks import TdmQf, TdmPolicy
from rlkit.torch.vae.vae_trainer import ConvVAETrainer
from rlkit.torch.online_vae.online_vae_tdm_td3 import OnlineVaeTdmTd3
preprocess_rl_variant(variant)
env = get_envs(variant)
es = get_exploration_strategy(variant, env)
vae_trainer_kwargs = variant.get('vae_trainer_kwargs')
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)
goal_dim = (env.observation_space.spaces[desired_goal_key].low.size)
action_dim = env.action_space.low.size
vectorized = 'vectorized' in env.reward_type
variant['algo_kwargs']['tdm_td3_kwargs']['tdm_kwargs'][
'vectorized'] = vectorized
norm_order = env.norm_order
# variant['algo_kwargs']['tdm_td3_kwargs']['tdm_kwargs'][
# 'norm_order'] = norm_order
qf1 = TdmQf(env=env,
vectorized=vectorized,
norm_order=norm_order,
observation_dim=obs_dim,
goal_dim=goal_dim,
action_dim=action_dim,
**variant['qf_kwargs'])
qf2 = TdmQf(env=env,
vectorized=vectorized,
norm_order=norm_order,
observation_dim=obs_dim,
goal_dim=goal_dim,
action_dim=action_dim,
**variant['qf_kwargs'])
policy = TdmPolicy(env=env,
observation_dim=obs_dim,
goal_dim=goal_dim,
action_dim=action_dim,
**variant['policy_kwargs'])
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
vae = env.vae
replay_buffer = OnlineVaeRelabelingBuffer(
vae=vae,
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']['tdm_td3_kwargs']
td3_kwargs = algo_kwargs['td3_kwargs']
td3_kwargs['training_env'] = env
tdm_kwargs = algo_kwargs['tdm_kwargs']
tdm_kwargs['observation_key'] = observation_key
tdm_kwargs['desired_goal_key'] = desired_goal_key
algo_kwargs["replay_buffer"] = replay_buffer
t = ConvVAETrainer(variant['vae_train_data'],
variant['vae_test_data'],
vae,
beta=variant['online_vae_beta'],
**vae_trainer_kwargs)
render = variant["render"]
assert 'vae_training_schedule' not in variant, "Just put it in algo_kwargs"
algorithm = OnlineVaeTdmTd3(
online_vae_kwargs=dict(vae=vae,
vae_trainer=t,
**variant['algo_kwargs']['online_vae_kwargs']),
tdm_td3_kwargs=dict(env=env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
**variant['algo_kwargs']['tdm_td3_kwargs']),
)
algorithm.to(ptu.device)
vae.to(ptu.device)
if variant.get("save_video", True):
policy.train(False)
rollout_function = rf.create_rollout_function(
rf.tdm_rollout,
init_tau=algorithm._sample_max_tau_for_rollout(),
decrement_tau=algorithm.cycle_taus_for_rollout,
cycle_tau=algorithm.cycle_taus_for_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,
algorithm.eval_policy,
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 tdm_td3_experiment(variant):
import rlkit.samplers.rollout_functions as rf
import rlkit.torch.pytorch_util as ptu
from rlkit.core import logger
from rlkit.data_management.obs_dict_replay_buffer import \
ObsDictRelabelingBuffer
from rlkit.exploration_strategies.base import (
PolicyWrappedWithExplorationStrategy)
from rlkit.state_distance.tdm_networks import TdmQf, TdmPolicy
from rlkit.state_distance.tdm_td3 import TdmTd3
preprocess_rl_variant(variant)
env = get_envs(variant)
es = get_exploration_strategy(variant, env)
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)
goal_dim = (env.observation_space.spaces[desired_goal_key].low.size)
action_dim = env.action_space.low.size
vectorized = 'vectorized' in env.reward_type
norm_order = env.norm_order
variant['algo_kwargs']['tdm_kwargs']['vectorized'] = vectorized
variant['qf_kwargs']['vectorized'] = vectorized
variant['qf_kwargs']['norm_order'] = norm_order
qf1 = TdmQf(env=env,
observation_dim=obs_dim,
goal_dim=goal_dim,
action_dim=action_dim,
**variant['qf_kwargs'])
qf2 = TdmQf(env=env,
observation_dim=obs_dim,
goal_dim=goal_dim,
action_dim=action_dim,
**variant['qf_kwargs'])
policy = TdmPolicy(env=env,
observation_dim=obs_dim,
goal_dim=goal_dim,
action_dim=action_dim,
**variant['policy_kwargs'])
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
variant['replay_buffer_kwargs']['vectorized'] = vectorized
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"]
tdm_kwargs = algo_kwargs['tdm_kwargs']
tdm_kwargs['observation_key'] = observation_key
tdm_kwargs['desired_goal_key'] = desired_goal_key
algorithm = TdmTd3(env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
**variant['algo_kwargs'])
algorithm.to(ptu.device)
if not variant.get("do_state_exp", False):
env.vae.to(ptu.device)
if variant.get("save_video", True):
logdir = logger.get_snapshot_dir()
policy.train(False)
rollout_function = rf.create_rollout_function(
rf.tdm_rollout,
init_tau=algorithm.max_tau,
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_train_funcs.append(video_func)
algorithm.train()
def tdm_td3_experiment(variant):
variant['env_kwargs'].update(variant['reward_params'])
env = variant['env_class'](**variant['env_kwargs'])
multiworld_env = variant.get('multiworld_env', True)
if multiworld_env is not True:
env = MultitaskEnvToSilentMultitaskEnv(env)
if variant["render"]:
env.pause_on_goal = True
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)
if multiworld_env is True:
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
else:
obs_dim = action_dim = goal_dim = None
vectorized = 'vectorized' in env.reward_type
variant['algo_kwargs']['tdm_kwargs']['vectorized'] = vectorized
norm_order = env.norm_order
variant['algo_kwargs']['tdm_kwargs']['norm_order'] = norm_order
qf1 = TdmQf(env=env,
observation_dim=obs_dim,
action_dim=action_dim,
goal_dim=goal_dim,
vectorized=vectorized,
norm_order=norm_order,
**variant['qf_kwargs'])
qf2 = TdmQf(env=env,
observation_dim=obs_dim,
action_dim=action_dim,
goal_dim=goal_dim,
vectorized=vectorized,
norm_order=norm_order,
**variant['qf_kwargs'])
policy = TdmPolicy(env=env,
observation_dim=obs_dim,
action_dim=action_dim,
goal_dim=goal_dim,
**variant['policy_kwargs'])
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
relabeling_env = pickle.loads(pickle.dumps(env))
algo_kwargs = variant['algo_kwargs']
if multiworld_env is True:
observation_key = variant.get('observation_key', 'state_observation')
desired_goal_key = variant.get('desired_goal_key',
'state_desired_goal')
achieved_goal_key = variant.get('achieved_goal_key',
'state_achieved_goal')
replay_buffer = ObsDictRelabelingBuffer(
env=relabeling_env,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
achieved_goal_key=achieved_goal_key,
vectorized=vectorized,
**variant['replay_buffer_kwargs'])
algo_kwargs['tdm_kwargs']['observation_key'] = observation_key
algo_kwargs['tdm_kwargs']['desired_goal_key'] = desired_goal_key
else:
replay_buffer = RelabelingReplayBuffer(
env=relabeling_env, **variant['replay_buffer_kwargs'])
# qf_criterion = variant['qf_criterion_class']()
# algo_kwargs['td3_kwargs']['qf_criterion'] = qf_criterion
algo_kwargs['td3_kwargs']['training_env'] = env
if 'tau_schedule_kwargs' in variant:
tau_schedule = IntPiecewiseLinearSchedule(
**variant['tau_schedule_kwargs'])
else:
tau_schedule = None
algo_kwargs['tdm_kwargs']['epoch_max_tau_schedule'] = tau_schedule
algorithm = TdmTd3(env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
replay_buffer=replay_buffer,
**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()