def __init__(
self,
env,
qf1,
qf2,
exploration_policy,
td3_kwargs,
tdm_kwargs,
base_kwargs,
policy=None,
eval_policy=None,
replay_buffer=None,
optimizer_class=optim.Adam,
):
TD3.__init__(self,
env=env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
replay_buffer=replay_buffer,
eval_policy=eval_policy,
optimizer_class=optimizer_class,
**td3_kwargs,
**base_kwargs)
super().__init__(**tdm_kwargs)
def __init__(
self,
*args,
her_kwargs,
td3_kwargs,
**kwargs
):
HER.__init__(self, **her_kwargs)
TD3.__init__(self, *args, **kwargs, **td3_kwargs)
assert isinstance(
self.replay_buffer, RelabelingReplayBuffer
) or isinstance(
self.replay_buffer, ObsDictRelabelingBuffer
)
def __init__(self,
*args,
observation_key='observation',
desired_goal_key='desired_goal',
**kwargs):
HER.__init__(
self,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
)
TD3.__init__(self, *args, **kwargs)
assert isinstance(self.replay_buffer,
RelabelingReplayBuffer) or isinstance(
self.replay_buffer, ObsDictRelabelingBuffer)
def experiment(variant):
env = NormalizedBoxEnv(variant['env_class']())
es = GaussianStrategy(action_space=env.action_space,
**variant['es_kwargs'])
obs_dim = env.observation_space.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'])
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
algorithm = TD3(env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
**variant['algo_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
if variant['multitask']:
env = CylinderXYPusher2DEnv(**variant['env_kwargs'])
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 = ConcatMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[400, 300],
)
qf2 = ConcatMlp(
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']
)
algorithm.to(ptu.device)
algorithm.train()
def td3_experiment(variant):
env = variant['env_class'](**variant['env_kwargs'])
env = MultitaskToFlatEnv(env)
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)
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 = 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']
)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
algorithm = TD3(
env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
**variant['algo_kwargs']
)
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
logger.add_text_output('./d_text.txt')
logger.add_tabular_output('./d_tabular.txt')
logger.set_snapshot_dir('./snaps')
farmer = Farmer([('0.0.0.0', 1)])
remote_env = farmer.force_acq_env()
remote_env.set_spaces()
env = NormalizedBoxEnv(remote_env)
es = GaussianStrategy(
action_space=env.action_space,
max_sigma=0.1,
min_sigma=0.1, # Constant sigma
)
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=[256, 256],
)
qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[256, 256],
)
policy = TanhMlpPolicy(
input_size=obs_dim,
output_size=action_dim,
hidden_sizes=[256, 256],
)
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 ptu.gpu_enabled():
algorithm.cuda()
algorithm.train()
def experiment(variant):
#Robot
env = gym.make('replab-v0')._start_rospy(goal_oriented=False)
#SIM
#env = gym.make('replab-v0')._start_sim(goal_oriented=False, render=False)
env.action_space.low *= 10
env.action_space.high *= 10
env = NormalizedBoxEnv(env)
es = GaussianStrategy(
action_space=env.action_space,
max_sigma=0.1,
min_sigma=0.1, # Constant sigma
)
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']
)
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
env = NormalizedBoxEnv(HopperEnv())
es = GaussianStrategy(
action_space=env.action_space,
max_sigma=0.1,
min_sigma=0.1, # Constant sigma
)
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 ptu.gpu_enabled():
algorithm.cuda()
algorithm.train()
def td3_experiment_offpolicy_online_vae(variant):
import rlkit.torch.pytorch_util as ptu
from rlkit.data_management.online_vae_replay_buffer import \
OnlineVaeRelabelingBuffer
from rlkit.torch.networks import ConcatMlp, TanhMlpPolicy
from rlkit.torch.vae.vae_trainer import ConvVAETrainer
from rlkit.torch.td3.td3 import TD3
from rlkit.exploration_strategies.base import (
PolicyWrappedWithExplorationStrategy)
from rlkit.exploration_strategies.gaussian_and_epislon import \
GaussianAndEpislonStrategy
from rlkit.torch.vae.online_vae_offpolicy_algorithm import OnlineVaeOffpolicyAlgorithm
preprocess_rl_variant(variant)
env = get_envs(variant)
uniform_dataset_fn = variant.get('generate_uniform_dataset_fn', None)
if uniform_dataset_fn:
uniform_dataset = uniform_dataset_fn(
**variant['generate_uniform_dataset_kwargs'])
else:
uniform_dataset = None
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,
)
target_qf1 = ConcatMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=hidden_sizes,
)
target_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,
# **variant['policy_kwargs']
)
target_policy = TanhMlpPolicy(
input_size=obs_dim,
output_size=action_dim,
hidden_sizes=hidden_sizes,
# **variant['policy_kwargs']
)
es = GaussianAndEpislonStrategy(
action_space=env.action_space,
max_sigma=.2,
min_sigma=.2, # constant sigma
epsilon=.3,
)
expl_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
vae = env.vae
replay_buffer_class = variant.get("replay_buffer_class",
OnlineVaeRelabelingBuffer)
replay_buffer = replay_buffer_class(vae=env.vae,
env=env,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
achieved_goal_key=achieved_goal_key,
**variant['replay_buffer_kwargs'])
replay_buffer.representation_size = vae.representation_size
vae_trainer_class = variant.get("vae_trainer_class", ConvVAETrainer)
vae_trainer = vae_trainer_class(env.vae,
**variant['online_vae_trainer_kwargs'])
assert 'vae_training_schedule' not in variant, "Just put it in algo_kwargs"
max_path_length = variant['max_path_length']
trainer = TD3(policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
target_policy=target_policy,
**variant['td3_trainer_kwargs'])
trainer = HERTrainer(trainer)
eval_path_collector = VAEWrappedEnvPathCollector(
variant['evaluation_goal_sampling_mode'],
env,
policy,
max_path_length,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
)
expl_path_collector = VAEWrappedEnvPathCollector(
variant['exploration_goal_sampling_mode'],
env,
expl_policy,
max_path_length,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
)
algorithm = OnlineVaeOffpolicyAlgorithm(
trainer=trainer,
exploration_env=env,
evaluation_env=env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
vae=vae,
vae_trainer=vae_trainer,
uniform_dataset=uniform_dataset,
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)
if variant['custom_goal_sampler'] == 'replay_buffer':
env.custom_goal_sampler = replay_buffer.sample_buffer_goals
algorithm.to(ptu.device)
vae.to(ptu.device)
algorithm.pretrain()
algorithm.train()
def td3_experiment_online_vae_exploring(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.torch.her.online_vae_joint_algo import OnlineVaeHerJointAlgo
from rlkit.torch.networks import ConcatMlp, TanhMlpPolicy
from rlkit.torch.td3.td3 import TD3
from rlkit.torch.vae.vae_trainer import ConvVAETrainer
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 +
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'],
)
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
exploring_qf1 = ConcatMlp(
input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs'],
)
exploring_qf2 = ConcatMlp(
input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs'],
)
exploring_policy = TanhMlpPolicy(
input_size=obs_dim,
output_size=action_dim,
**variant['policy_kwargs'],
)
exploring_exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=exploring_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'])
variant["algo_kwargs"]["replay_buffer"] = replay_buffer
if variant.get('use_replay_buffer_goals', False):
env.replay_buffer = replay_buffer
env.use_replay_buffer_goals = True
vae_trainer_kwargs = variant.get('vae_trainer_kwargs')
t = ConvVAETrainer(variant['vae_train_data'],
variant['vae_test_data'],
vae,
beta=variant['online_vae_beta'],
**vae_trainer_kwargs)
control_algorithm = TD3(env=env,
training_env=env,
qf1=qf1,
qf2=qf2,
policy=policy,
exploration_policy=exploration_policy,
**variant['algo_kwargs'])
exploring_algorithm = TD3(env=env,
training_env=env,
qf1=exploring_qf1,
qf2=exploring_qf2,
policy=exploring_policy,
exploration_policy=exploring_exploration_policy,
**variant['algo_kwargs'])
assert 'vae_training_schedule' not in variant,\
"Just put it in joint_algo_kwargs"
algorithm = OnlineVaeHerJointAlgo(vae=vae,
vae_trainer=t,
env=env,
training_env=env,
policy=policy,
exploration_policy=exploration_policy,
replay_buffer=replay_buffer,
algo1=control_algorithm,
algo2=exploring_algorithm,
algo1_prefix="Control_",
algo2_prefix="VAE_Exploration_",
observation_key=observation_key,
desired_goal_key=desired_goal_key,
**variant['joint_algo_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.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,
algorithm.eval_policy,
variant,
)
algorithm.post_train_funcs.append(video_func)
algorithm.train()
def experiment(variant):
expl_env = variant['env_class'](**variant['env_kwargs'])
eval_env = variant['env_class'](**variant['env_kwargs'])
observation_key = 'state_observation'
desired_goal_key = 'state_desired_goal'
achieved_goal_key = desired_goal_key.replace("desired", "achieved")
es = GaussianAndEpislonStrategy(
action_space=expl_env.action_space,
max_sigma=.2,
min_sigma=.2, # constant sigma
epsilon=.3,
)
obs_dim = expl_env.observation_space.spaces['observation'].low.size
goal_dim = expl_env.observation_space.spaces['desired_goal'].low.size
action_dim = expl_env.action_space.low.size
qf1 = ConcatMlp(input_size=obs_dim + goal_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
qf2 = ConcatMlp(input_size=obs_dim + goal_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
target_qf1 = ConcatMlp(input_size=obs_dim + goal_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
target_qf2 = ConcatMlp(input_size=obs_dim + goal_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
policy = TanhMlpPolicy(input_size=obs_dim + goal_dim,
output_size=action_dim,
**variant['policy_kwargs'])
target_policy = TanhMlpPolicy(input_size=obs_dim + goal_dim,
output_size=action_dim,
**variant['policy_kwargs'])
expl_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
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'])
trainer = TD3(policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
target_policy=target_policy,
**variant['trainer_kwargs'])
trainer = HERTrainer(trainer)
eval_path_collector = GoalConditionedPathCollector(
eval_env,
policy,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
)
expl_path_collector = GoalConditionedPathCollector(
expl_env,
expl_policy,
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'])
algorithm.to(ptu.device)
algorithm.train()
def 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_state_experiment_video_save_function
from rlkit.torch.her.her_td3 import HerTd3
from rlkit.torch.td3.td3 import TD3
from rlkit.torch.networks import ConcatMlp, TanhMlpPolicy
from rlkit.data_management.obs_dict_replay_buffer import (
ObsDictReplayBuffer)
from rlkit.torch.torch_rl_algorithm import TorchBatchRLAlgorithm
from rlkit.samplers.data_collector.path_collector import ObsDictPathCollector
if 'env_id' in variant:
eval_env = gym.make(variant['env_id'])
expl_env = gym.make(variant['env_id'])
else:
eval_env_kwargs = variant.get('eval_env_kwargs', variant['env_kwargs'])
eval_env = variant['env_class'](**eval_env_kwargs)
expl_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 = ObsDictReplayBuffer(
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
exploration_type = variant['exploration_type']
if exploration_type == 'ou':
es = OUStrategy(action_space=eval_env.action_space,
**variant['es_kwargs'])
elif exploration_type == 'gaussian':
es = GaussianStrategy(
action_space=eval_env.action_space,
**variant['es_kwargs'],
)
elif exploration_type == 'epsilon':
es = EpsilonGreedy(
action_space=eval_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'])
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'])
target_policy = TanhMlpPolicy(input_size=obs_dim + goal_dim,
output_size=action_dim,
**variant['policy_kwargs'])
expl_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
trainer = TD3(policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
target_policy=target_policy,
**variant['trainer_kwargs'])
observation_key = 'observation'
desired_goal_key = 'desired_goal'
eval_path_collector = ObsDictPathCollector(
eval_env,
policy,
observation_key=observation_key,
# render=True,
# desired_goal_key=desired_goal_key,
)
expl_path_collector = ObsDictPathCollector(
expl_env,
expl_policy,
observation_key=observation_key,
# render=True,
# 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=observation_key,
# desired_goal_key=algorithm.desired_goal_key,
# )
# video_func = get_state_experiment_video_save_function(
# rollout_function,
# env,
# policy,
# variant,
# )
# algorithm.post_epoch_funcs.append(video_func)
algorithm.to(ptu.device)
algorithm.train()
def her_td3_experiment(variant):
import gym
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.exploration_strategies.gaussian_and_epislon import \
GaussianAndEpislonStrategy
from rlkit.launchers.launcher_util import setup_logger
from rlkit.samplers.data_collector import GoalConditionedPathCollector
from rlkit.torch.her.her import HERTrainer
from rlkit.torch.networks import ConcatMlp, TanhMlpPolicy
from rlkit.torch.td3.td3 import TD3
from rlkit.torch.torch_rl_algorithm import TorchBatchRLAlgorithm
import rlkit.samplers.rollout_functions as rf
from rlkit.torch.grill.launcher import get_state_experiment_video_save_function
if 'env_id' in variant:
eval_env = gym.make(variant['env_id'])
expl_env = gym.make(variant['env_id'])
else:
eval_env_kwargs = variant.get('eval_env_kwargs', variant['env_kwargs'])
eval_env = variant['env_class'](**eval_env_kwargs)
expl_env = variant['env_class'](**variant['env_kwargs'])
observation_key = 'state_observation'
desired_goal_key = 'state_desired_goal'
achieved_goal_key = desired_goal_key.replace("desired", "achieved")
es = GaussianAndEpislonStrategy(
action_space=expl_env.action_space,
max_sigma=.2,
min_sigma=.2, # constant sigma
epsilon=.3,
)
obs_dim = expl_env.observation_space.spaces['observation'].low.size
goal_dim = expl_env.observation_space.spaces['desired_goal'].low.size
action_dim = expl_env.action_space.low.size
qf1 = ConcatMlp(input_size=obs_dim + goal_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
qf2 = ConcatMlp(input_size=obs_dim + goal_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
target_qf1 = ConcatMlp(input_size=obs_dim + goal_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
target_qf2 = ConcatMlp(input_size=obs_dim + goal_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
policy = TanhMlpPolicy(input_size=obs_dim + goal_dim,
output_size=action_dim,
**variant['policy_kwargs'])
target_policy = TanhMlpPolicy(input_size=obs_dim + goal_dim,
output_size=action_dim,
**variant['policy_kwargs'])
expl_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
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'])
trainer = TD3(policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
target_policy=target_policy,
**variant['trainer_kwargs'])
trainer = HERTrainer(trainer)
eval_path_collector = GoalConditionedPathCollector(
eval_env,
policy,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
)
expl_path_collector = GoalConditionedPathCollector(
expl_env,
expl_policy,
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=observation_key,
desired_goal_key=desired_goal_key,
)
video_func = get_state_experiment_video_save_function(
rollout_function,
eval_env,
policy,
variant,
)
algorithm.post_epoch_funcs.append(video_func)
algorithm.to(ptu.device)
algorithm.train()
def state_td3bc_experiment(variant):
if variant.get('env_id', None):
import gym
import multiworld
multiworld.register_all_envs()
eval_env = gym.make(variant['env_id'])
eval_env = MujocoGymToMultiEnv(eval_env)
# eval_env = EncoderWrappedEnv(eval_env)
expl_env = gym.make(variant['env_id'])
expl_env = MujocoGymToMultiEnv(expl_env)
# expl_env = EncoderWrappedEnv(expl_env)
else:
eval_env_kwargs = variant.get('eval_env_kwargs', variant['env_kwargs'])
eval_env = variant['env_class'](**eval_env_kwargs)
expl_env = variant['env_class'](**variant['env_kwargs'])
observation_key = 'state_observation'
desired_goal_key = 'state_desired_goal'
achieved_goal_key = desired_goal_key.replace("desired", "achieved")
es_strat = variant.get('es', 'ou')
if es_strat == 'ou':
es = OUStrategy(
action_space=expl_env.action_space,
max_sigma=variant['exploration_noise'],
min_sigma=variant['exploration_noise'],
)
elif es_strat == 'gauss_eps':
es = GaussianAndEpislonStrategy(
action_space=expl_env.action_space,
max_sigma=variant['exploration_noise'],
min_sigma=variant['exploration_noise'], # constant sigma
epsilon=0,
)
else:
raise ValueError("invalid exploration strategy provided")
obs_dim = expl_env.observation_space.spaces['observation'].low.size
goal_dim = 0 # expl_env.observation_space.spaces['desired_goal'].low.size
action_dim = expl_env.action_space.low.size
qf1 = ConcatMlp(input_size=obs_dim + goal_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
qf2 = ConcatMlp(input_size=obs_dim + goal_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
target_qf1 = ConcatMlp(input_size=obs_dim + goal_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
target_qf2 = ConcatMlp(input_size=obs_dim + goal_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
policy = TanhMlpPolicy(input_size=obs_dim + goal_dim,
output_size=action_dim,
**variant['policy_kwargs'])
target_policy = TanhMlpPolicy(input_size=obs_dim + goal_dim,
output_size=action_dim,
**variant['policy_kwargs'])
expl_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
replay_buffer = ObsDictReplayBuffer(
env=eval_env,
observation_key=observation_key,
# desired_goal_key=desired_goal_key,
# achieved_goal_key=achieved_goal_key,
**variant['replay_buffer_kwargs'])
demo_train_buffer = ObsDictReplayBuffer(
env=eval_env,
observation_key=observation_key,
# desired_goal_key=desired_goal_key,
# achieved_goal_key=achieved_goal_key,
max_size=variant['replay_buffer_kwargs']['max_size'])
demo_test_buffer = ObsDictReplayBuffer(
env=eval_env,
observation_key=observation_key,
# desired_goal_key=desired_goal_key,
# achieved_goal_key=achieved_goal_key,
max_size=variant['replay_buffer_kwargs']['max_size'],
)
if variant.get('td3_bc', True):
td3_trainer = TD3BCTrainer(env=expl_env,
policy=policy,
qf1=qf1,
qf2=qf2,
replay_buffer=replay_buffer,
demo_train_buffer=demo_train_buffer,
demo_test_buffer=demo_test_buffer,
target_qf1=target_qf1,
target_qf2=target_qf2,
target_policy=target_policy,
**variant['td3_bc_trainer_kwargs'])
else:
td3_trainer = TD3(policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
target_policy=target_policy,
**variant['td3_trainer_kwargs'])
trainer = td3_trainer # HERTrainer(td3_trainer)
eval_path_collector = ObsDictPathCollector( # GoalConditionedPathCollector(
eval_env,
policy,
observation_key=observation_key,
# desired_goal_key=desired_goal_key,
)
expl_path_collector = ObsDictPathCollector( # GoalConditionedPathCollector(
expl_env,
expl_policy,
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", True):
if variant.get("presampled_goals", None):
variant['image_env_kwargs'][
'presampled_goals'] = load_local_or_remote_file(
variant['presampled_goals']).item()
image_eval_env = ImageEnv(eval_env, **variant["image_env_kwargs"])
image_eval_path_collector = ObsDictPathCollector( # GoalConditionedPathCollector(
image_eval_env,
policy,
observation_key='state_observation',
# desired_goal_key='state_desired_goal',
)
image_expl_env = ImageEnv(expl_env, **variant["image_env_kwargs"])
image_expl_path_collector = ObsDictPathCollector( # GoalConditionedPathCollector(
image_expl_env,
expl_policy,
observation_key='state_observation',
# desired_goal_key='state_desired_goal',
)
video_func = VideoSaveFunction(
image_eval_env,
variant,
image_expl_path_collector,
image_eval_path_collector,
)
algorithm.post_train_funcs.append(video_func)
algorithm.to(ptu.device)
if variant.get('load_demos', False):
td3_trainer.load_demos()
if variant.get('pretrain_policy', False):
td3_trainer.pretrain_policy_with_bc()
if variant.get('pretrain_rl', False):
td3_trainer.pretrain_q_with_bc_data()
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 rlkit.envs.wrappers import ImageMujocoEnv, NormalizedBoxEnv
from rlkit.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 = ConcatMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[400, 300],
)
qf2 = ConcatMlp(
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)