def experiment(variant):
imsize = variant['imsize']
env = ImageForkReacher2dEnv(variant["arm_goal_distance_cost_coeff"],
variant["arm_object_distance_cost_coeff"],
[imsize, imsize, 3],
goal_object_distance_cost_coeff=variant[
"goal_object_distance_cost_coeff"],
ctrl_cost_coeff=variant["ctrl_cost_coeff"])
partial_obs_size = env.obs_dim - imsize * imsize * 3
print("partial dim was " + str(partial_obs_size))
env = NormalizedBoxEnv(env)
obs_dim = int(np.prod(env.observation_space.shape))
action_dim = int(np.prod(env.action_space.shape))
qf1 = MergedCNN(input_width=imsize,
input_height=imsize,
output_size=1,
input_channels=3,
added_fc_input_size=action_dim,
**variant['cnn_params'])
qf2 = MergedCNN(input_width=imsize,
input_height=imsize,
output_size=1,
input_channels=3,
added_fc_input_size=action_dim,
**variant['cnn_params'])
vf = CNN(input_width=imsize,
input_height=imsize,
output_size=1,
input_channels=3,
**variant['cnn_params'])
policy = TanhCNNGaussianPolicy(input_width=imsize,
input_height=imsize,
output_size=action_dim,
input_channels=3,
**variant['cnn_params'])
algorithm = SACTrainer(env=env,
policy=policy,
qf1=qf1,
qf2=qf2,
vf=vf,
**variant['algo_params'])
algorithm.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 = FlattenMlp(input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs'])
qf2 = FlattenMlp(input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs'])
target_qf1 = FlattenMlp(input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs'])
target_qf2 = FlattenMlp(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()
def _disentangled_her_twin_sac_experiment_v2(
max_path_length,
encoder_kwargs,
disentangled_qf_kwargs,
qf_kwargs,
twin_sac_trainer_kwargs,
replay_buffer_kwargs,
policy_kwargs,
evaluation_goal_sampling_mode,
exploration_goal_sampling_mode,
algo_kwargs,
save_video=True,
env_id=None,
env_class=None,
env_kwargs=None,
observation_key='state_observation',
desired_goal_key='state_desired_goal',
achieved_goal_key='state_achieved_goal',
# Video parameters
latent_dim=2,
save_video_kwargs=None,
**kwargs
):
import railrl.samplers.rollout_functions as rf
import railrl.torch.pytorch_util as ptu
from railrl.data_management.obs_dict_replay_buffer import \
ObsDictRelabelingBuffer
from railrl.torch.networks import FlattenMlp
from railrl.torch.sac.policies import TanhGaussianPolicy
from railrl.torch.torch_rl_algorithm import TorchBatchRLAlgorithm
if save_video_kwargs is None:
save_video_kwargs = {}
if env_kwargs is None:
env_kwargs = {}
assert env_id or env_class
if env_id:
import gym
import multiworld
multiworld.register_all_envs()
train_env = gym.make(env_id)
eval_env = gym.make(env_id)
else:
eval_env = env_class(**env_kwargs)
train_env = env_class(**env_kwargs)
obs_dim = train_env.observation_space.spaces[observation_key].low.size
goal_dim = train_env.observation_space.spaces[desired_goal_key].low.size
action_dim = train_env.action_space.low.size
encoder = FlattenMlp(
input_size=goal_dim,
output_size=latent_dim,
**encoder_kwargs
)
qf1 = DisentangledMlpQf(
goal_processor=encoder,
preprocess_obs_dim=obs_dim,
action_dim=action_dim,
qf_kwargs=qf_kwargs,
**disentangled_qf_kwargs
)
qf2 = DisentangledMlpQf(
goal_processor=encoder,
preprocess_obs_dim=obs_dim,
action_dim=action_dim,
qf_kwargs=qf_kwargs,
**disentangled_qf_kwargs
)
target_qf1 = DisentangledMlpQf(
goal_processor=Detach(encoder),
preprocess_obs_dim=obs_dim,
action_dim=action_dim,
qf_kwargs=qf_kwargs,
**disentangled_qf_kwargs
)
target_qf2 = DisentangledMlpQf(
goal_processor=Detach(encoder),
preprocess_obs_dim=obs_dim,
action_dim=action_dim,
qf_kwargs=qf_kwargs,
**disentangled_qf_kwargs
)
policy = TanhGaussianPolicy(
obs_dim=obs_dim + goal_dim,
action_dim=action_dim,
**policy_kwargs
)
replay_buffer = ObsDictRelabelingBuffer(
env=train_env,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
achieved_goal_key=achieved_goal_key,
**replay_buffer_kwargs
)
sac_trainer = SACTrainer(
env=train_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**twin_sac_trainer_kwargs
)
trainer = HERTrainer(sac_trainer)
eval_path_collector = GoalConditionedPathCollector(
eval_env,
MakeDeterministic(policy),
max_path_length,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
goal_sampling_mode=evaluation_goal_sampling_mode,
)
expl_path_collector = GoalConditionedPathCollector(
train_env,
policy,
max_path_length,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
goal_sampling_mode=exploration_goal_sampling_mode,
)
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=train_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
max_path_length=max_path_length,
**algo_kwargs,
)
algorithm.to(ptu.device)
if save_video:
save_vf_heatmap = save_video_kwargs.get('save_vf_heatmap', True)
def v_function(obs):
action = policy.get_actions(obs)
obs, action = ptu.from_numpy(obs), ptu.from_numpy(action)
return qf1(obs, action, return_individual_q_vals=True)
add_heatmap = partial(add_heatmap_imgs_to_o_dict, v_function=v_function)
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,
full_o_postprocess_func=add_heatmap if save_vf_heatmap else None,
)
img_keys = ['v_vals'] + [
'v_vals_dim_{}'.format(dim) for dim
in range(latent_dim)
]
eval_video_func = get_save_video_function(
rollout_function,
eval_env,
MakeDeterministic(policy),
tag="eval",
get_extra_imgs=partial(get_extra_imgs, img_keys=img_keys),
**save_video_kwargs
)
train_video_func = get_save_video_function(
rollout_function,
train_env,
policy,
tag="train",
get_extra_imgs=partial(get_extra_imgs, img_keys=img_keys),
**save_video_kwargs
)
decoder = FlattenMlp(
input_size=obs_dim,
output_size=obs_dim,
hidden_sizes=[128, 128],
)
decoder.to(ptu.device)
# algorithm.post_train_funcs.append(train_decoder(variant, encoder, decoder))
# algorithm.post_train_funcs.append(plot_encoder_function(variant, encoder))
# algorithm.post_train_funcs.append(plot_buffer_function(
# save_video_period, 'state_achieved_goal'))
# algorithm.post_train_funcs.append(plot_buffer_function(
# save_video_period, 'state_desired_goal'))
algorithm.post_train_funcs.append(eval_video_func)
algorithm.post_train_funcs.append(train_video_func)
algorithm.train()
def experiment(variant):
import multiworld.envs.pygame
env = gym.make('Point2DEnv-ImageFixedGoal-v0')
input_width, input_height = env.image_shape
action_dim = int(np.prod(env.action_space.shape))
cnn_params = variant['cnn_params']
cnn_params.update(
input_width=input_width,
input_height=input_height,
input_channels=3,
output_conv_channels=True,
output_size=None,
)
if variant['shared_qf_conv']:
qf_cnn = PretrainedCNN(**cnn_params)
qf1 = MlpQfWithObsProcessor(
obs_processor=nn.Sequential(qf_cnn, Flatten()),
output_size=1,
input_size=action_dim + qf_cnn.conv_output_flat_size,
**variant['qf_kwargs'])
qf2 = MlpQfWithObsProcessor(
obs_processor=nn.Sequential(qf_cnn, Flatten()),
output_size=1,
input_size=action_dim + qf_cnn.conv_output_flat_size,
**variant['qf_kwargs'])
target_qf_cnn = PretrainedCNN(**cnn_params)
target_qf1 = MlpQfWithObsProcessor(
obs_processor=nn.Sequential(target_qf_cnn, Flatten()),
output_size=1,
input_size=action_dim + target_qf_cnn.conv_output_flat_size,
**variant['qf_kwargs'])
target_qf2 = MlpQfWithObsProcessor(
obs_processor=nn.Sequential(target_qf_cnn, Flatten()),
output_size=1,
input_size=action_dim + target_qf_cnn.conv_output_flat_size,
**variant['qf_kwargs'])
else:
qf1_cnn = PretrainedCNN(**cnn_params)
cnn_output_dim = qf1_cnn.conv_output_flat_size
qf1 = MlpQfWithObsProcessor(obs_processor=nn.Sequential(
qf1_cnn, Flatten()),
output_size=1,
input_size=action_dim + cnn_output_dim,
**variant['qf_kwargs'])
qf2 = MlpQfWithObsProcessor(obs_processor=nn.Sequential(
PretrainedCNN(**cnn_params), Flatten()),
output_size=1,
input_size=action_dim + cnn_output_dim,
**variant['qf_kwargs'])
target_qf1 = MlpQfWithObsProcessor(
obs_processor=nn.Sequential(PretrainedCNN(**cnn_params),
Flatten()),
output_size=1,
input_size=action_dim + cnn_output_dim,
**variant['qf_kwargs'])
target_qf2 = MlpQfWithObsProcessor(
obs_processor=nn.Sequential(PretrainedCNN(**cnn_params),
Flatten()),
output_size=1,
input_size=action_dim + cnn_output_dim,
**variant['qf_kwargs'])
action_dim = int(np.prod(env.action_space.shape))
policy_cnn = PretrainedCNN(**cnn_params)
policy = TanhGaussianPolicyAdapter(nn.Sequential(policy_cnn, Flatten()),
policy_cnn.conv_output_flat_size,
action_dim, **variant['policy_kwargs'])
eval_env = expl_env = env
eval_policy = MakeDeterministic(policy)
eval_path_collector = MdpPathCollector(
eval_env, eval_policy, **variant['eval_path_collector_kwargs'])
replay_buffer = EnvReplayBuffer(
variant['replay_buffer_size'],
expl_env,
)
trainer = SACTrainer(env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**variant['trainer_kwargs'])
if variant['collection_mode'] == 'batch':
expl_path_collector = MdpPathCollector(
expl_env, policy, **variant['expl_path_collector_kwargs'])
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'])
elif variant['collection_mode'] == 'online':
expl_path_collector = MdpStepCollector(
expl_env, policy, **variant['expl_path_collector_kwargs'])
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,
**variant['algo_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def _pointmass_fixed_goal_experiment(vae_latent_size,
replay_buffer_size,
cnn_kwargs,
vae_kwargs,
policy_kwargs,
qf_kwargs,
e2e_trainer_kwargs,
sac_trainer_kwargs,
algorithm_kwargs,
eval_path_collector_kwargs=None,
expl_path_collector_kwargs=None,
**kwargs):
if expl_path_collector_kwargs is None:
expl_path_collector_kwargs = {}
if eval_path_collector_kwargs is None:
eval_path_collector_kwargs = {}
from multiworld.core.image_env import ImageEnv
from multiworld.envs.pygame.point2d import Point2DEnv
from multiworld.core.flat_goal_env import FlatGoalEnv
env = Point2DEnv(
images_are_rgb=True,
render_onscreen=False,
show_goal=False,
ball_radius=2,
render_size=48,
fixed_goal=(0, 0),
)
env = ImageEnv(env, imsize=env.render_size, transpose=True, normalize=True)
env = FlatGoalEnv(env) #, append_goal_to_obs=True)
input_width, input_height = env.image_shape
action_dim = int(np.prod(env.action_space.shape))
vae = ConvVAE(
representation_size=vae_latent_size,
input_channels=3,
imsize=input_width,
decoder_output_activation=nn.Sigmoid(),
# decoder_distribution='gaussian_identity_variance',
**vae_kwargs)
encoder = Vae2Encoder(vae)
def make_cnn():
return networks.CNN(input_width=input_width,
input_height=input_height,
input_channels=3,
output_conv_channels=True,
output_size=None,
**cnn_kwargs)
def make_qf():
return networks.MlpQfWithObsProcessor(obs_processor=nn.Sequential(
encoder,
networks.Flatten(),
),
output_size=1,
input_size=action_dim +
vae_latent_size,
**qf_kwargs)
qf1 = make_qf()
qf2 = make_qf()
target_qf1 = make_qf()
target_qf2 = make_qf()
action_dim = int(np.prod(env.action_space.shape))
policy_cnn = make_cnn()
policy = TanhGaussianPolicyAdapter(
nn.Sequential(policy_cnn, networks.Flatten()),
policy_cnn.conv_output_flat_size, action_dim, **policy_kwargs)
eval_env = expl_env = env
eval_policy = MakeDeterministic(policy)
eval_path_collector = MdpPathCollector(eval_env, eval_policy,
**eval_path_collector_kwargs)
replay_buffer = EnvReplayBuffer(
replay_buffer_size,
expl_env,
)
vae_trainer = VAETrainer(vae)
sac_trainer = SACTrainer(env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**sac_trainer_kwargs)
trainer = End2EndSACTrainer(
sac_trainer=sac_trainer,
vae_trainer=vae_trainer,
**e2e_trainer_kwargs,
)
expl_path_collector = MdpPathCollector(expl_env, policy,
**expl_path_collector_kwargs)
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,
**algorithm_kwargs)
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
env = RobosuiteStateWrapperEnv(wrapped_env_id=variant['env_id'],
**variant['env_kwargs']) #
obs_dim = env.observation_space.low.size
action_dim = env.action_space.low.size
hidden_sizes = variant['hidden_sizes']
qf1 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=hidden_sizes,
)
qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=hidden_sizes,
)
target_qf1 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=hidden_sizes,
)
target_qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=hidden_sizes,
)
policy = TanhGaussianPolicy(
obs_dim=obs_dim,
action_dim=action_dim,
hidden_sizes=hidden_sizes,
)
es = OUStrategy(action_space=env.action_space,
max_sigma=variant['exploration_noise'],
min_sigma=variant['exploration_noise'])
exploration_policy = PolicyWrappedWithExplorationStrategy(
exploration_strategy=es,
policy=policy,
)
eval_policy = MakeDeterministic(policy)
eval_path_collector = MdpPathCollector(
env,
eval_policy,
)
expl_path_collector = MdpPathCollector(
env,
exploration_policy,
)
replay_buffer = EnvReplayBuffer(
variant['replay_buffer_size'],
env,
)
trainer = SACTrainer(env=env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**variant['trainer_kwargs'])
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=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)
algorithm.train()
def experiment(variant):
expl_env = roboverse.make(variant['env'],
gui=False,
randomize=variant['randomize_env'],
observation_mode=variant['obs'],
reward_type='shaped',
transpose_image=True)
if variant['obs'] == 'pixels_debug':
robot_state_dims = 11
elif variant['obs'] == 'pixels':
robot_state_dims = 4
else:
raise NotImplementedError
expl_env = FlatEnv(expl_env,
use_robot_state=variant['use_robot_state'],
robot_state_dims=robot_state_dims)
eval_env = expl_env
img_width, img_height = eval_env.image_shape
num_channels = 3
action_dim = int(np.prod(eval_env.action_space.shape))
cnn_params = variant['cnn_params']
cnn_params.update(
input_width=img_width,
input_height=img_height,
input_channels=num_channels,
)
if variant['use_robot_state']:
cnn_params.update(
added_fc_input_size=robot_state_dims,
output_conv_channels=False,
hidden_sizes=[400, 400],
output_size=200,
)
else:
cnn_params.update(
added_fc_input_size=0,
output_conv_channels=True,
output_size=None,
)
qf_cnn = CNN(**cnn_params)
if variant['use_robot_state']:
qf_obs_processor = qf_cnn
else:
qf_obs_processor = nn.Sequential(
qf_cnn,
Flatten(),
)
qf_kwargs = copy.deepcopy(variant['qf_kwargs'])
qf_kwargs['obs_processor'] = qf_obs_processor
qf_kwargs['output_size'] = 1
if variant['use_robot_state']:
qf_kwargs['input_size'] = (action_dim + qf_cnn.output_size)
else:
qf_kwargs['input_size'] = (action_dim + qf_cnn.conv_output_flat_size)
qf1 = MlpQfWithObsProcessor(**qf_kwargs)
qf2 = MlpQfWithObsProcessor(**qf_kwargs)
target_qf_cnn = CNN(**cnn_params)
if variant['use_robot_state']:
target_qf_obs_processor = target_qf_cnn
else:
target_qf_obs_processor = nn.Sequential(
target_qf_cnn,
Flatten(),
)
target_qf_kwargs = copy.deepcopy(variant['qf_kwargs'])
target_qf_kwargs['obs_processor'] = target_qf_obs_processor
target_qf_kwargs['output_size'] = 1
if variant['use_robot_state']:
target_qf_kwargs['input_size'] = (action_dim +
target_qf_cnn.output_size)
else:
target_qf_kwargs['input_size'] = (action_dim +
target_qf_cnn.conv_output_flat_size)
target_qf1 = MlpQfWithObsProcessor(**target_qf_kwargs)
target_qf2 = MlpQfWithObsProcessor(**target_qf_kwargs)
action_dim = int(np.prod(eval_env.action_space.shape))
policy_cnn = CNN(**cnn_params)
if variant['use_robot_state']:
policy_obs_processor = policy_cnn
else:
policy_obs_processor = nn.Sequential(
policy_cnn,
Flatten(),
)
if variant['use_robot_state']:
policy = TanhGaussianPolicyAdapter(policy_obs_processor,
policy_cnn.output_size, action_dim,
**variant['policy_kwargs'])
else:
policy = TanhGaussianPolicyAdapter(policy_obs_processor,
policy_cnn.conv_output_flat_size,
action_dim,
**variant['policy_kwargs'])
eval_policy = MakeDeterministic(policy)
eval_path_collector = MdpPathCollector(
eval_env, eval_policy, **variant['eval_path_collector_kwargs'])
replay_buffer = EnvReplayBuffer(
variant['replay_buffer_size'],
expl_env,
)
trainer = SACTrainer(env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**variant['trainer_kwargs'])
if variant['collection_mode'] == 'batch':
expl_path_collector = MdpPathCollector(
expl_env, policy, **variant['expl_path_collector_kwargs'])
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'])
elif variant['collection_mode'] == 'online':
expl_path_collector = MdpStepCollector(
expl_env, policy, **variant['expl_path_collector_kwargs'])
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,
**variant['algo_kwargs'])
else:
raise NotImplementedError
video_func = VideoSaveFunctionBullet(variant)
algorithm.post_train_funcs.append(video_func)
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
expl_env = PointmassBaseEnv()
eval_env = expl_env
action_dim = int(np.prod(eval_env.action_space.shape))
state_dim = 3
qf_kwargs = copy.deepcopy(variant['qf_kwargs'])
qf_kwargs['output_size'] = 1
qf_kwargs['input_size'] = action_dim + state_dim
qf1 = MlpQf(**qf_kwargs)
qf2 = MlpQf(**qf_kwargs)
target_qf_kwargs = copy.deepcopy(qf_kwargs)
target_qf1 = MlpQf(**target_qf_kwargs)
target_qf2 = MlpQf(**target_qf_kwargs)
policy_kwargs = copy.deepcopy(variant['policy_kwargs'])
policy_kwargs['action_dim'] = action_dim
policy_kwargs['obs_dim'] = state_dim
policy = TanhGaussianPolicy(**policy_kwargs)
eval_policy = MakeDeterministic(policy)
eval_path_collector = MdpPathCollector(
eval_env, eval_policy, **variant['eval_path_collector_kwargs'])
replay_buffer = EnvReplayBuffer(
variant['replay_buffer_size'],
expl_env,
)
trainer = SACTrainer(env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**variant['trainer_kwargs'])
if variant['collection_mode'] == 'batch':
expl_path_collector = MdpPathCollector(
expl_env, policy, **variant['expl_path_collector_kwargs'])
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'])
elif variant['collection_mode'] == 'online':
expl_path_collector = MdpStepCollector(
expl_env, policy, **variant['expl_path_collector_kwargs'])
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,
**variant['algo_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
import multiworld
multiworld.register_all_envs()
env = gym.make('Image48SawyerReachXYEnv-v1')
observation_key = 'image_proprio_observation'
input_width, input_height = env.image_shape
action_dim = int(np.prod(env.action_space.shape))
cnn_params = variant['cnn_params']
cnn_params.update(
input_width=input_width,
input_height=input_height,
input_channels=3,
added_fc_input_size=3,
output_conv_channels=True,
output_size=None,
)
if variant['shared_qf_conv']:
qf_cnn = CNN(**cnn_params)
qf1 = MlpQfWithObsProcessor(
obs_processor=nn.Sequential(qf_cnn, Flatten()),
output_size=1,
input_size=action_dim + qf_cnn.conv_output_flat_size,
**variant['qf_kwargs'])
qf2 = MlpQfWithObsProcessor(
obs_processor=nn.Sequential(qf_cnn, Flatten()),
output_size=1,
input_size=action_dim + qf_cnn.conv_output_flat_size,
**variant['qf_kwargs'])
target_qf_cnn = CNN(**cnn_params)
target_qf1 = MlpQfWithObsProcessor(
obs_processor=nn.Sequential(target_qf_cnn, Flatten()),
output_size=1,
input_size=action_dim + target_qf_cnn.conv_output_flat_size,
**variant['qf_kwargs'])
target_qf2 = MlpQfWithObsProcessor(
obs_processor=nn.Sequential(target_qf_cnn, Flatten()),
output_size=1,
input_size=action_dim + target_qf_cnn.conv_output_flat_size,
**variant['qf_kwargs'])
else:
qf1_cnn = CNN(**cnn_params)
cnn_output_dim = qf1_cnn.conv_output_flat_size
qf1 = MlpQfWithObsProcessor(obs_processor=nn.Sequential(
qf1_cnn, Flatten()),
output_size=1,
input_size=action_dim + cnn_output_dim,
**variant['qf_kwargs'])
qf2 = MlpQfWithObsProcessor(obs_processor=nn.Sequential(
CNN(**cnn_params), Flatten()),
output_size=1,
input_size=action_dim + cnn_output_dim,
**variant['qf_kwargs'])
target_qf1 = MlpQfWithObsProcessor(
obs_processor=nn.Sequential(CNN(**cnn_params), Flatten()),
output_size=1,
input_size=action_dim + cnn_output_dim,
**variant['qf_kwargs'])
target_qf2 = MlpQfWithObsProcessor(
obs_processor=nn.Sequential(CNN(**cnn_params), Flatten()),
output_size=1,
input_size=action_dim + cnn_output_dim,
**variant['qf_kwargs'])
policy_cnn = CNN(**cnn_params)
policy = TanhGaussianPolicyAdapter(nn.Sequential(policy_cnn, Flatten()),
policy_cnn.conv_output_flat_size,
action_dim, **variant['policy_kwargs'])
eval_env = expl_env = env
eval_policy = MakeDeterministic(policy)
eval_path_collector = ObsDictPathCollector(
eval_env,
eval_policy,
observation_key=observation_key,
**variant['eval_path_collector_kwargs'])
replay_buffer = ObsDictReplayBuffer(
variant['replay_buffer_size'],
expl_env,
observation_key=observation_key,
**variant['replay_buffer_kwargs'],
)
trainer = SACTrainer(env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**variant['trainer_kwargs'])
if variant['collection_mode'] == 'batch':
expl_path_collector = ObsDictPathCollector(
expl_env,
policy,
observation_key=observation_key,
**variant['expl_path_collector_kwargs'])
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'])
elif variant['collection_mode'] == 'online':
expl_path_collector = ObsDictStepCollector(
expl_env,
policy,
observation_key=observation_key,
**variant['expl_path_collector_kwargs'])
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,
**variant['algo_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def active_representation_learning_experiment(variant):
import railrl.torch.pytorch_util as ptu
from railrl.data_management.obs_dict_replay_buffer import ObsDictReplayBuffer
from railrl.torch.networks import FlattenMlp
from railrl.torch.sac.policies import TanhGaussianPolicy
from railrl.torch.arl.active_representation_learning_algorithm import \
ActiveRepresentationLearningAlgorithm
from railrl.torch.arl.representation_wrappers import RepresentationWrappedEnv
from multiworld.core.image_env import ImageEnv
from railrl.samplers.data_collector import MdpPathCollector
preprocess_rl_variant(variant)
model_class = variant.get('model_class')
model_kwargs = variant.get('model_kwargs')
model = model_class(**model_kwargs)
model.representation_size = 4
model.imsize = 48
variant["vae_path"] = model
reward_params = variant.get("reward_params", dict())
init_camera = variant.get("init_camera", None)
env = variant["env_class"](**variant['env_kwargs'])
image_env = ImageEnv(
env,
variant.get('imsize'),
init_camera=init_camera,
transpose=True,
normalize=True,
)
env = RepresentationWrappedEnv(
image_env,
model,
)
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
action_dim = env.action_space.low.size
hidden_sizes = variant.get('hidden_sizes', [400, 300])
qf1 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=hidden_sizes,
)
qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=hidden_sizes,
)
target_qf1 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=hidden_sizes,
)
target_qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=hidden_sizes,
)
policy = TanhGaussianPolicy(
obs_dim=obs_dim,
action_dim=action_dim,
hidden_sizes=hidden_sizes,
)
vae = env.vae
replay_buffer = ObsDictReplayBuffer(env=env,
**variant['replay_buffer_kwargs'])
model_trainer_class = variant.get('model_trainer_class')
model_trainer_kwargs = variant.get('model_trainer_kwargs')
model_trainer = model_trainer_class(
model,
**model_trainer_kwargs,
)
# vae_trainer = ConvVAETrainer(
# 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 = 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)
eval_path_collector = MdpPathCollector(
env,
MakeDeterministic(policy),
# max_path_length,
# observation_key=observation_key,
# desired_goal_key=desired_goal_key,
)
expl_path_collector = MdpPathCollector(
env,
policy,
# max_path_length,
# observation_key=observation_key,
# desired_goal_key=desired_goal_key,
)
algorithm = ActiveRepresentationLearningAlgorithm(
trainer=trainer,
exploration_env=env,
evaluation_env=env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
model=model,
model_trainer=model_trainer,
uniform_dataset=uniform_dataset,
max_path_length=max_path_length,
**variant['algo_kwargs'])
algorithm.to(ptu.device)
vae.to(ptu.device)
algorithm.train()
def experiment(variant):
expl_env = roboverse.make(variant['env'],
gui=False,
randomize=variant['randomize_env'],
observation_mode=variant['obs'],
reward_type='shaped',
transpose_image=True)
eval_env = expl_env
img_width, img_height = eval_env.image_shape
num_channels = 3
action_dim = int(np.prod(eval_env.action_space.shape))
cnn_params = variant['cnn_params']
cnn_params.update(
input_width=img_width,
input_height=img_height,
input_channels=num_channels,
added_fc_input_size=0,
output_conv_channels=True,
output_size=None,
)
qf_cnn = CNN(**cnn_params)
qf_obs_processor = nn.Sequential(
qf_cnn,
Flatten(),
)
qf_kwargs = copy.deepcopy(variant['qf_kwargs'])
qf_kwargs['obs_processor'] = qf_obs_processor
qf_kwargs['output_size'] = 1
qf_kwargs['input_size'] = (action_dim + qf_cnn.conv_output_flat_size)
qf1 = MlpQfWithObsProcessor(**qf_kwargs)
qf2 = MlpQfWithObsProcessor(**qf_kwargs)
target_qf_cnn = CNN(**cnn_params)
target_qf_obs_processor = nn.Sequential(
target_qf_cnn,
Flatten(),
)
target_qf_kwargs = copy.deepcopy(variant['qf_kwargs'])
target_qf_kwargs['obs_processor'] = target_qf_obs_processor
target_qf_kwargs['output_size'] = 1
target_qf_kwargs['input_size'] = (action_dim +
target_qf_cnn.conv_output_flat_size)
target_qf1 = MlpQfWithObsProcessor(**target_qf_kwargs)
target_qf2 = MlpQfWithObsProcessor(**target_qf_kwargs)
action_dim = int(np.prod(eval_env.action_space.shape))
policy_cnn = CNN(**cnn_params)
policy_obs_processor = nn.Sequential(
policy_cnn,
Flatten(),
)
policy = TanhGaussianPolicyAdapter(policy_obs_processor,
policy_cnn.conv_output_flat_size,
action_dim, **variant['policy_kwargs'])
observation_key = 'image'
eval_policy = MakeDeterministic(policy)
eval_path_collector = ObsDictPathCollector(
eval_env,
eval_policy,
observation_key=observation_key,
**variant['eval_path_collector_kwargs'])
with open(variant['buffer'], 'rb') as f:
replay_buffer = pickle.load(f)
trainer = SACTrainer(env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**variant['trainer_kwargs'])
expl_path_collector = ObsDictStepCollector(
expl_env,
policy,
observation_key=observation_key,
**variant['expl_path_collector_kwargs'])
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,
q_learning_alg=False, # should be false for SAC
batch_rl=variant['batch_rl'],
**variant['algo_kwargs'])
video_func = VideoSaveFunctionBullet(variant)
# dump_buffer_func = BufferSaveFunction(variant)
algorithm.post_train_funcs.append(video_func)
# algorithm.post_train_funcs.append(dump_buffer_func)
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
expl_env = roboverse.make(variant['env'], gui=False, randomize=True,
observation_mode='state', reward_type='shaped',
transpose_image=True)
eval_env = expl_env
action_dim = int(np.prod(eval_env.action_space.shape))
state_dim = 11
qf_kwargs = copy.deepcopy(variant['qf_kwargs'])
qf_kwargs['output_size'] = 1
qf_kwargs['input_size'] = action_dim + state_dim
qf1 = MlpQf(**qf_kwargs)
qf2 = MlpQf(**qf_kwargs)
target_qf_kwargs = copy.deepcopy(qf_kwargs)
target_qf1 = MlpQf(**target_qf_kwargs)
target_qf2 = MlpQf(**target_qf_kwargs)
policy_kwargs = copy.deepcopy(variant['policy_kwargs'])
policy_kwargs['action_dim'] = action_dim
policy_kwargs['obs_dim'] = state_dim
policy = TanhGaussianPolicy(**policy_kwargs)
eval_policy = MakeDeterministic(policy)
eval_path_collector = MdpPathCollector(
eval_env,
eval_policy,
**variant['eval_path_collector_kwargs']
)
with open(variant['buffer'], 'rb') as f:
replay_buffer = pickle.load(f)
trainer = SACTrainer(
env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**variant['trainer_kwargs']
)
expl_path_collector = MdpPathCollector(
expl_env,
policy,
**variant['expl_path_collector_kwargs']
)
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,
q_learning_alg=False, # should be false for SAC
batch_rl=variant['batch_rl'],
**variant['algo_kwargs']
)
dump_buffer_func = BufferSaveFunction(variant)
# algorithm.post_train_funcs.append(dump_buffer_func)
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
env_params = dict(
block_random=0.3,
camera_random=0,
simple_observations=False,
continuous=True,
remove_height_hack=True,
render_mode="DIRECT",
# render_mode="GUI",
num_objects=5,
max_num_training_models=900,
target=False,
test=False,
)
expl_env = FlatEnv(KukaGraspingProceduralEnv(**env_params))
eval_env = expl_env
img_width, img_height = eval_env.image_shape
num_channels = 3
action_dim = int(np.prod(eval_env.action_space.shape))
cnn_params = variant['cnn_params']
cnn_params.update(
input_width=img_width,
input_height=img_height,
input_channels=num_channels,
added_fc_input_size=0,
output_conv_channels=True,
output_size=None,
)
qf_cnn = CNN(**cnn_params)
qf_obs_processor = nn.Sequential(
qf_cnn,
Flatten(),
)
qf_kwargs = copy.deepcopy(variant['qf_kwargs'])
qf_kwargs['obs_processor'] = qf_obs_processor
qf_kwargs['output_size'] = 1
qf_kwargs['input_size'] = (action_dim + qf_cnn.conv_output_flat_size)
qf1 = MlpQfWithObsProcessor(**qf_kwargs)
qf2 = MlpQfWithObsProcessor(**qf_kwargs)
target_qf_cnn = CNN(**cnn_params)
target_qf_obs_processor = nn.Sequential(
target_qf_cnn,
Flatten(),
)
target_qf_kwargs = copy.deepcopy(variant['qf_kwargs'])
target_qf_kwargs['obs_processor'] = target_qf_obs_processor
target_qf_kwargs['output_size'] = 1
target_qf_kwargs['input_size'] = (action_dim +
target_qf_cnn.conv_output_flat_size)
target_qf1 = MlpQfWithObsProcessor(**target_qf_kwargs)
target_qf2 = MlpQfWithObsProcessor(**target_qf_kwargs)
action_dim = int(np.prod(eval_env.action_space.shape))
policy_cnn = CNN(**cnn_params)
policy_obs_processor = nn.Sequential(
policy_cnn,
Flatten(),
)
policy = TanhGaussianPolicyAdapter(policy_obs_processor,
policy_cnn.conv_output_flat_size,
action_dim, **variant['policy_kwargs'])
observation_key = 'image'
eval_policy = MakeDeterministic(policy)
eval_path_collector = ObsDictPathCollector(
eval_env,
eval_policy,
observation_key=observation_key,
**variant['eval_path_collector_kwargs'])
replay_buffer = ObsDictReplayBuffer(
variant['replay_buffer_size'],
expl_env,
observation_key=observation_key,
)
trainer = SACTrainer(env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**variant['trainer_kwargs'])
if variant['collection_mode'] == 'batch':
expl_path_collector = ObsDictPathCollector(
expl_env,
policy,
observation_key=observation_key,
**variant['expl_path_collector_kwargs'])
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'])
elif variant['collection_mode'] == 'online':
expl_path_collector = ObsDictStepCollector(
expl_env,
policy,
observation_key=observation_key,
**variant['expl_path_collector_kwargs'])
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,
**variant['algo_kwargs'])
else:
raise NotImplementedError
video_func = VideoSaveFunctionBullet(variant)
algorithm.post_train_funcs.append(video_func)
# dump_buffer_func = BufferSaveFunction(variant)
# algorithm.post_train_funcs.append(dump_buffer_func)
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
import gym
from multiworld.envs.mujoco import register_custom_envs
from multiworld.core.flat_goal_env import FlatGoalEnv
register_custom_envs()
expl_env = FlatGoalEnv(
gym.make(variant['env_id']),
obs_keys=['state_observation'],
goal_keys=['xy_desired_goal'],
append_goal_to_obs=False,
)
eval_env = FlatGoalEnv(
gym.make(variant['env_id']),
obs_keys=['state_observation'],
goal_keys=['xy_desired_goal'],
append_goal_to_obs=False,
)
obs_dim = expl_env.observation_space.low.size
action_dim = eval_env.action_space.low.size
M = variant['layer_size']
qf1 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[M, M],
)
qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[M, M],
)
target_qf1 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[M, M],
)
target_qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[M, M],
)
policy = TanhGaussianPolicy(
obs_dim=obs_dim,
action_dim=action_dim,
hidden_sizes=[M, M],
)
eval_policy = MakeDeterministic(policy)
eval_path_collector = MdpPathCollector(
eval_env,
eval_policy,
)
replay_buffer = EnvReplayBuffer(
variant['replay_buffer_size'],
expl_env,
)
trainer = SACTrainer(
env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**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:
expl_path_collector = MdpPathCollector(
expl_env,
policy,
)
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)
algorithm.train()
def her_sac_experiment(
max_path_length,
qf_kwargs,
twin_sac_trainer_kwargs,
replay_buffer_kwargs,
policy_kwargs,
evaluation_goal_sampling_mode,
exploration_goal_sampling_mode,
algo_kwargs,
save_video=True,
env_id=None,
env_class=None,
env_kwargs=None,
observation_key='state_observation',
desired_goal_key='state_desired_goal',
achieved_goal_key='state_achieved_goal',
# Video parameters
save_video_kwargs=None,
**kwargs
):
import railrl.samplers.rollout_functions as rf
import railrl.torch.pytorch_util as ptu
from railrl.data_management.obs_dict_replay_buffer import \
ObsDictRelabelingBuffer
from railrl.torch.networks import FlattenMlp
from railrl.torch.sac.policies import TanhGaussianPolicy
from railrl.torch.torch_rl_algorithm import TorchBatchRLAlgorithm
if not save_video_kwargs:
save_video_kwargs = {}
if env_kwargs is None:
env_kwargs = {}
assert env_id or env_class
if env_id:
import gym
import multiworld
multiworld.register_all_envs()
train_env = gym.make(env_id)
eval_env = gym.make(env_id)
else:
eval_env = env_class(**env_kwargs)
train_env = env_class(**env_kwargs)
obs_dim = (
train_env.observation_space.spaces[observation_key].low.size
+ train_env.observation_space.spaces[desired_goal_key].low.size
)
action_dim = train_env.action_space.low.size
qf1 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
**qf_kwargs
)
qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
**qf_kwargs
)
target_qf1 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
**qf_kwargs
)
target_qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
**qf_kwargs
)
policy = TanhGaussianPolicy(
obs_dim=obs_dim,
action_dim=action_dim,
**policy_kwargs
)
replay_buffer = ObsDictRelabelingBuffer(
env=train_env,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
achieved_goal_key=achieved_goal_key,
**replay_buffer_kwargs
)
trainer = SACTrainer(
env=train_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**twin_sac_trainer_kwargs
)
trainer = HERTrainer(trainer)
eval_path_collector = GoalConditionedPathCollector(
eval_env,
MakeDeterministic(policy),
max_path_length,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
goal_sampling_mode=evaluation_goal_sampling_mode,
)
expl_path_collector = GoalConditionedPathCollector(
train_env,
policy,
max_path_length,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
goal_sampling_mode=exploration_goal_sampling_mode,
)
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=train_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
max_path_length=max_path_length,
**algo_kwargs
)
algorithm.to(ptu.device)
if save_video:
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,
return_dict_obs=True,
)
eval_video_func = get_save_video_function(
rollout_function,
eval_env,
MakeDeterministic(policy),
tag="eval",
**save_video_kwargs
)
train_video_func = get_save_video_function(
rollout_function,
train_env,
policy,
tag="train",
**save_video_kwargs
)
# algorithm.post_train_funcs.append(plot_buffer_function(
# save_video_period, 'state_achieved_goal'))
# algorithm.post_train_funcs.append(plot_buffer_function(
# save_video_period, 'state_desired_goal'))
algorithm.post_train_funcs.append(eval_video_func)
algorithm.post_train_funcs.append(train_video_func)
algorithm.train()
def experiment(variant):
# from softlearning.environments.gym import register_image_reach
# register_image_reach()
# env = gym.envs.make(
# 'Pusher2d-ImageReach-v0',
# )
from softlearning.environments.gym.mujoco.image_pusher_2d import (
ImageForkReacher2dEnv)
env_kwargs = {
'image_shape': (32, 32, 3),
'arm_goal_distance_cost_coeff': 1.0,
'arm_object_distance_cost_coeff': 0.0,
}
eval_env = ImageForkReacher2dEnv(**env_kwargs)
expl_env = ImageForkReacher2dEnv(**env_kwargs)
input_width, input_height, input_channels = eval_env.image_shape
image_dim = input_width * input_height * input_channels
action_dim = int(np.prod(eval_env.action_space.shape))
cnn_params = variant['cnn_params']
cnn_params.update(
input_width=input_width,
input_height=input_height,
input_channels=input_channels,
added_fc_input_size=4,
output_conv_channels=True,
output_size=None,
)
non_image_dim = int(np.prod(eval_env.observation_space.shape)) - image_dim
if variant['shared_qf_conv']:
qf_cnn = CNN(**cnn_params)
qf_obs_processor = nn.Sequential(
Split(qf_cnn, identity, image_dim),
FlattenEach(),
Concat(),
)
qf_kwargs = copy.deepcopy(variant['qf_kwargs'])
qf_kwargs['obs_processor'] = qf_obs_processor
qf_kwargs['output_size'] = 1
qf_kwargs['input_size'] = (action_dim + qf_cnn.conv_output_flat_size +
non_image_dim)
qf1 = MlpQfWithObsProcessor(**qf_kwargs)
qf2 = MlpQfWithObsProcessor(**qf_kwargs)
target_qf_cnn = CNN(**cnn_params)
target_qf_obs_processor = nn.Sequential(
Split(target_qf_cnn, identity, image_dim),
FlattenEach(),
Concat(),
)
target_qf_kwargs = copy.deepcopy(variant['qf_kwargs'])
target_qf_kwargs['obs_processor'] = target_qf_obs_processor
target_qf_kwargs['output_size'] = 1
target_qf_kwargs['input_size'] = (action_dim +
target_qf_cnn.conv_output_flat_size +
non_image_dim)
target_qf1 = MlpQfWithObsProcessor(**target_qf_kwargs)
target_qf2 = MlpQfWithObsProcessor(**target_qf_kwargs)
else:
qf1_cnn = CNN(**cnn_params)
cnn_output_dim = qf1_cnn.conv_output_flat_size
qf1 = MlpQfWithObsProcessor(obs_processor=qf1_cnn,
output_size=1,
input_size=action_dim + cnn_output_dim,
**variant['qf_kwargs'])
qf2 = MlpQfWithObsProcessor(obs_processor=CNN(**cnn_params),
output_size=1,
input_size=action_dim + cnn_output_dim,
**variant['qf_kwargs'])
target_qf1 = MlpQfWithObsProcessor(obs_processor=CNN(**cnn_params),
output_size=1,
input_size=action_dim +
cnn_output_dim,
**variant['qf_kwargs'])
target_qf2 = MlpQfWithObsProcessor(obs_processor=CNN(**cnn_params),
output_size=1,
input_size=action_dim +
cnn_output_dim,
**variant['qf_kwargs'])
action_dim = int(np.prod(eval_env.action_space.shape))
policy_cnn = CNN(**cnn_params)
policy_obs_processor = nn.Sequential(
Split(policy_cnn, identity, image_dim),
FlattenEach(),
Concat(),
)
policy = TanhGaussianPolicyAdapter(
policy_obs_processor, policy_cnn.conv_output_flat_size + non_image_dim,
action_dim, **variant['policy_kwargs'])
eval_policy = MakeDeterministic(policy)
eval_path_collector = MdpPathCollector(
eval_env, eval_policy, **variant['eval_path_collector_kwargs'])
replay_buffer = EnvReplayBuffer(
variant['replay_buffer_size'],
expl_env,
)
trainer = SACTrainer(env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**variant['trainer_kwargs'])
if variant['collection_mode'] == 'batch':
expl_path_collector = MdpPathCollector(
expl_env, policy, **variant['expl_path_collector_kwargs'])
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'])
elif variant['collection_mode'] == 'online':
expl_path_collector = MdpStepCollector(
expl_env, policy, **variant['expl_path_collector_kwargs'])
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,
**variant['algo_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
import gym
from multiworld.envs.mujoco import register_custom_envs
register_custom_envs()
observation_key = 'state_observation'
desired_goal_key = 'xy_desired_goal'
expl_env = gym.make(variant['env_id'])
eval_env = gym.make(variant['env_id'])
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
M = variant['layer_size']
qf1 = FlattenMlp(
input_size=obs_dim + action_dim + goal_dim,
output_size=1,
hidden_sizes=[M, M],
)
qf2 = FlattenMlp(
input_size=obs_dim + action_dim + goal_dim,
output_size=1,
hidden_sizes=[M, M],
)
target_qf1 = FlattenMlp(
input_size=obs_dim + action_dim + goal_dim,
output_size=1,
hidden_sizes=[M, M],
)
target_qf2 = FlattenMlp(
input_size=obs_dim + action_dim + goal_dim,
output_size=1,
hidden_sizes=[M, M],
)
policy = TanhGaussianPolicy(
obs_dim=obs_dim + goal_dim,
action_dim=action_dim,
hidden_sizes=[M, M],
)
eval_policy = MakeDeterministic(policy)
eval_path_collector = GoalConditionedPathCollector(
eval_env,
eval_policy,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
)
trainer = SACTrainer(env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**variant['trainer_kwargs'])
trainer = HERTrainer(trainer)
if variant['collection_mode'] == 'online':
expl_step_collector = GoalConditionedStepCollector(
expl_env,
policy,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
)
algorithm = TorchOnlineRLAlgorithm(
trainer=trainer,
exploration_env=expl_env,
evaluation_env=eval_env,
exploration_data_collector=expl_step_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
**variant['algo_kwargs'])
else:
expl_path_collector = GoalConditionedPathCollector(
expl_env,
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 experiment(variant):
from multiworld.envs.mujoco import register_goal_example_envs
register_goal_example_envs()
eval_env = gym.make('Image48SawyerPushForwardEnv-v0')
expl_env = gym.make('Image48SawyerPushForwardEnv-v0')
# Hack for now
eval_env.wrapped_env.transpose = True
expl_env.wrapped_env.transpose = True
# More hacks, use a dense reward instead
eval_env.wrapped_env.wrapped_env.reward_type = 'puck_distance'
expl_env.wrapped_env.wrapped_env.reward_type = 'puck_distance'
img_width, img_height = eval_env.image_shape
num_channels = 3
action_dim = int(np.prod(eval_env.action_space.shape))
cnn_params = variant['cnn_params']
cnn_params.update(
input_width=img_width,
input_height=img_height,
input_channels=num_channels,
added_fc_input_size=4,
output_conv_channels=True,
output_size=None,
)
qf_cnn = CNN(**cnn_params)
qf_obs_processor = nn.Sequential(
qf_cnn,
Flatten(),
)
qf_kwargs = copy.deepcopy(variant['qf_kwargs'])
qf_kwargs['obs_processor'] = qf_obs_processor
qf_kwargs['output_size'] = 1
qf_kwargs['input_size'] = (action_dim + qf_cnn.conv_output_flat_size)
qf1 = MlpQfWithObsProcessor(**qf_kwargs)
qf2 = MlpQfWithObsProcessor(**qf_kwargs)
target_qf_cnn = CNN(**cnn_params)
target_qf_obs_processor = nn.Sequential(
target_qf_cnn,
Flatten(),
)
target_qf_kwargs = copy.deepcopy(variant['qf_kwargs'])
target_qf_kwargs['obs_processor'] = target_qf_obs_processor
target_qf_kwargs['output_size'] = 1
target_qf_kwargs['input_size'] = (action_dim +
target_qf_cnn.conv_output_flat_size)
target_qf1 = MlpQfWithObsProcessor(**target_qf_kwargs)
target_qf2 = MlpQfWithObsProcessor(**target_qf_kwargs)
action_dim = int(np.prod(eval_env.action_space.shape))
policy_cnn = CNN(**cnn_params)
policy_obs_processor = nn.Sequential(
policy_cnn,
Flatten(),
)
policy = TanhGaussianPolicyAdapter(policy_obs_processor,
policy_cnn.conv_output_flat_size,
action_dim, **variant['policy_kwargs'])
eval_policy = MakeDeterministic(policy)
eval_path_collector = MdpPathCollector(
eval_env, eval_policy, **variant['eval_path_collector_kwargs'])
replay_buffer = EnvReplayBuffer(
variant['replay_buffer_size'],
expl_env,
)
trainer = SACTrainer(env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**variant['trainer_kwargs'])
if variant['collection_mode'] == 'batch':
expl_path_collector = MdpPathCollector(
expl_env, policy, **variant['expl_path_collector_kwargs'])
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'])
elif variant['collection_mode'] == 'online':
expl_path_collector = MdpStepCollector(
expl_env, policy, **variant['expl_path_collector_kwargs'])
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,
**variant['algo_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def goal_conditioned_sac_experiment(
max_path_length,
qf_kwargs,
sac_trainer_kwargs,
replay_buffer_kwargs,
policy_kwargs,
algo_kwargs,
env_id=None,
env_class=None,
env_kwargs=None,
observation_key='state_observation',
desired_goal_key='state_desired_goal',
achieved_goal_key='state_achieved_goal',
contextual_env_kwargs=None,
evaluation_goal_sampling_mode=None,
exploration_goal_sampling_mode=None,
# Video parameters
save_video=True,
save_video_kwargs=None,
):
if contextual_env_kwargs is None:
contextual_env_kwargs = {}
if not save_video_kwargs:
save_video_kwargs = {}
def contextual_env_distrib_and_reward(env_id, env_class, env_kwargs,
goal_sampling_mode):
env = get_gym_env(env_id, env_class=env_class, env_kwargs=env_kwargs)
env.goal_sampling_mode = goal_sampling_mode
goal_distribution = GoalDistributionFromMultitaskEnv(
env,
desired_goal_key=desired_goal_key,
)
reward_fn = ContextualRewardFnFromMultitaskEnv(
env=env,
desired_goal_key=desired_goal_key,
achieved_goal_key=achieved_goal_key,
)
env = ContextualEnv(
env,
context_distribution=goal_distribution,
reward_fn=reward_fn,
observation_key=observation_key,
**contextual_env_kwargs,
)
return env, goal_distribution, reward_fn
expl_env, expl_context_distrib, expl_reward = contextual_env_distrib_and_reward(
env_id, env_class, env_kwargs, exploration_goal_sampling_mode)
eval_env, eval_context_distrib, eval_reward = contextual_env_distrib_and_reward(
env_id, env_class, env_kwargs, evaluation_goal_sampling_mode)
context_key = desired_goal_key
obs_dim = (expl_env.observation_space.spaces[observation_key].low.size +
expl_env.observation_space.spaces[desired_goal_key].low.size)
action_dim = expl_env.action_space.low.size
def create_qf():
return FlattenMlp(input_size=obs_dim + action_dim,
output_size=1,
**qf_kwargs)
qf1 = create_qf()
qf2 = create_qf()
target_qf1 = create_qf()
target_qf2 = create_qf()
policy = TanhGaussianPolicy(obs_dim=obs_dim,
action_dim=action_dim,
**policy_kwargs)
ob_keys_to_save = [
observation_key,
desired_goal_key,
achieved_goal_key,
]
def concat_context_to_obs(batch):
obs = batch['observations']
next_obs = batch['next_observations']
context = batch['contexts']
batch['observations'] = np.concatenate([obs, context], axis=1)
batch['next_observations'] = np.concatenate([next_obs, context],
axis=1)
return batch
replay_buffer = ContextualRelabelingReplayBuffer(
env=eval_env,
context_key=desired_goal_key,
context_distribution=eval_context_distrib,
sample_context_from_obs_dict_fn=SelectKeyFn(achieved_goal_key),
ob_keys_to_save=ob_keys_to_save,
reward_fn=eval_reward,
post_process_batch_fn=concat_context_to_obs,
**replay_buffer_kwargs)
trainer = SACTrainer(env=expl_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**sac_trainer_kwargs)
eval_path_collector = ContextualPathCollector(
eval_env,
MakeDeterministic(policy),
observation_key=observation_key,
context_key=context_key,
)
expl_path_collector = ContextualPathCollector(
expl_env,
policy,
observation_key=observation_key,
context_key=context_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,
max_path_length=max_path_length,
**algo_kwargs)
algorithm.to(ptu.device)
if save_video:
rollout_function = partial(
rf.contextual_rollout,
max_path_length=max_path_length,
observation_key=observation_key,
context_key=context_key,
)
eval_video_func = get_save_video_function(rollout_function,
eval_env,
MakeDeterministic(policy),
tag="eval",
**save_video_kwargs)
train_video_func = get_save_video_function(rollout_function,
expl_env,
policy,
tag="train",
**save_video_kwargs)
algorithm.post_train_funcs.append(eval_video_func)
algorithm.post_train_funcs.append(train_video_func)
algorithm.train()
def experiment(variant):
ptu.set_gpu_mode(True, 0)
from softlearning.environments.gym import register_image_reach
register_image_reach()
env = gym.make('Pusher2d-ImageReach-v0', arm_goal_distance_cost_coeff=1.0, arm_object_distance_cost_coeff=0.0)
#import ipdb; ipdb.set_trace()
input_width, input_height = env.image_shape
action_dim = int(np.prod(env.action_space.shape))
cnn_params = variant['cnn_params']
cnn_params.update(
input_width=input_width,
input_height=input_height,
input_channels=3,
added_fc_input_size=4,
output_conv_channels=True,
output_size=None,
)
if variant['shared_qf_conv']:
qf_cnn = CNN(**cnn_params)
qf1 = MlpQfWithObsProcessor(
obs_processor=qf_cnn,
output_size=1,
input_size=action_dim+qf_cnn.conv_output_flat_size,
**variant['qf_kwargs']
)
qf2 = MlpQfWithObsProcessor(
obs_processor=qf_cnn,
output_size=1,
input_size=action_dim+qf_cnn.conv_output_flat_size,
**variant['qf_kwargs']
)
target_qf_cnn = CNN(**cnn_params)
target_qf1 = MlpQfWithObsProcessor(
obs_processor=target_qf_cnn,
output_size=1,
input_size=action_dim+qf_cnn.conv_output_flat_size,
**variant['qf_kwargs']
)
target_qf2 = MlpQfWithObsProcessor(
obs_processor=target_qf_cnn,
output_size=1,
input_size=action_dim+qf_cnn.conv_output_flat_size,
**variant['qf_kwargs']
)
else:
qf1_cnn = CNN(**cnn_params)
cnn_output_dim = qf1_cnn.conv_output_flat_size
qf1 = MlpQfWithObsProcessor(
obs_processor=qf1_cnn,
output_size=1,
input_size=action_dim+cnn_output_dim,
**variant['qf_kwargs']
)
qf2 = MlpQfWithObsProcessor(
obs_processor=CNN(**cnn_params),
output_size=1,
input_size=action_dim+cnn_output_dim,
**variant['qf_kwargs']
)
target_qf1 = MlpQfWithObsProcessor(
obs_processor=CNN(**cnn_params),
output_size=1,
input_size=action_dim+cnn_output_dim,
**variant['qf_kwargs']
)
target_qf2 = MlpQfWithObsProcessor(
obs_processor=CNN(**cnn_params),
output_size=1,
input_size=action_dim+cnn_output_dim,
**variant['qf_kwargs']
)
action_dim = int(np.prod(env.action_space.shape))
policy_cnn = CNN(**cnn_params)
policy = TanhGaussianPolicyAdapter(
policy_cnn,
policy_cnn.conv_output_flat_size,
action_dim,
)
eval_env = expl_env = env
eval_policy = MakeDeterministic(policy)
eval_path_collector = MdpPathCollector(
eval_env,
eval_policy,
**variant['eval_path_collector_kwargs']
)
replay_buffer = EnvReplayBuffer(
variant['replay_buffer_size'],
expl_env,
)
trainer = SACTrainer(
env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**variant['trainer_kwargs']
)
if variant['collection_mode'] == 'batch':
expl_path_collector = MdpPathCollector(
expl_env,
policy,
**variant['expl_path_collector_kwargs']
)
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']
)
elif variant['collection_mode'] == 'online':
expl_path_collector = MdpStepCollector(
expl_env,
policy,
**variant['expl_path_collector_kwargs']
)
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,
**variant['algo_kwargs']
)
elif variant['collection_mode'] == 'parallel':
expl_path_collector = MdpPathCollector(
expl_env,
policy,
**variant['expl_path_collector_kwargs']
)
algorithm = TorchParallelRLAlgorithm(
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 twin_sac_experiment(variant):
import railrl.torch.pytorch_util as ptu
from railrl.data_management.obs_dict_replay_buffer import \
ObsDictRelabelingBuffer
from railrl.torch.networks import FlattenMlp
from railrl.torch.sac.policies import TanhGaussianPolicy
from railrl.torch.torch_rl_algorithm import TorchBatchRLAlgorithm
from railrl.torch.sac.policies import MakeDeterministic
from railrl.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 = FlattenMlp(input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
qf2 = FlattenMlp(input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
target_qf1 = FlattenMlp(input_size=obs_dim + action_dim,
output_size=1,
**variant['qf_kwargs'])
target_qf2 = FlattenMlp(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 grill_her_twin_sac_experiment_online_vae(variant):
import railrl.torch.pytorch_util as ptu
from railrl.data_management.online_vae_replay_buffer import \
OnlineVaeRelabelingBuffer
from railrl.torch.networks import FlattenMlp
from railrl.torch.sac.policies import TanhGaussianPolicy
from railrl.torch.vae.vae_trainer import ConvVAETrainer
grill_preprocess_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 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=hidden_sizes,
)
qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=hidden_sizes,
)
target_qf1 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=hidden_sizes,
)
target_qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=hidden_sizes,
)
policy = TanhGaussianPolicy(
obs_dim=obs_dim,
action_dim=action_dim,
hidden_sizes=hidden_sizes,
)
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'])
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 = 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)
eval_path_collector = VAEWrappedEnvPathCollector(
env,
MakeDeterministic(policy),
max_path_length,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
goal_sampling_mode=variant['evaluation_goal_sampling_mode'],
)
expl_path_collector = VAEWrappedEnvPathCollector(
env,
policy,
max_path_length,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
goal_sampling_mode=variant['exploration_goal_sampling_mode'],
)
algorithm = OnlineVaeAlgorithm(
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.train()
def _disentangled_grill_her_twin_sac_experiment(
max_path_length,
encoder_kwargs,
disentangled_qf_kwargs,
qf_kwargs,
twin_sac_trainer_kwargs,
replay_buffer_kwargs,
policy_kwargs,
vae_evaluation_goal_sampling_mode,
vae_exploration_goal_sampling_mode,
base_env_evaluation_goal_sampling_mode,
base_env_exploration_goal_sampling_mode,
algo_kwargs,
env_id=None,
env_class=None,
env_kwargs=None,
observation_key='state_observation',
desired_goal_key='state_desired_goal',
achieved_goal_key='state_achieved_goal',
latent_dim=2,
vae_wrapped_env_kwargs=None,
vae_path=None,
vae_n_vae_training_kwargs=None,
vectorized=False,
save_video=True,
save_video_kwargs=None,
have_no_disentangled_encoder=False,
**kwargs):
if env_kwargs is None:
env_kwargs = {}
assert env_id or env_class
if env_id:
import gym
import multiworld
multiworld.register_all_envs()
train_env = gym.make(env_id)
eval_env = gym.make(env_id)
else:
eval_env = env_class(**env_kwargs)
train_env = env_class(**env_kwargs)
train_env.goal_sampling_mode = base_env_exploration_goal_sampling_mode
eval_env.goal_sampling_mode = base_env_evaluation_goal_sampling_mode
if vae_path:
vae = load_local_or_remote_file(vae_path)
else:
vae = get_n_train_vae(latent_dim=latent_dim,
env=eval_env,
**vae_n_vae_training_kwargs)
train_env = VAEWrappedEnv(train_env,
vae,
imsize=train_env.imsize,
**vae_wrapped_env_kwargs)
eval_env = VAEWrappedEnv(eval_env,
vae,
imsize=train_env.imsize,
**vae_wrapped_env_kwargs)
obs_dim = train_env.observation_space.spaces[observation_key].low.size
goal_dim = train_env.observation_space.spaces[desired_goal_key].low.size
action_dim = train_env.action_space.low.size
encoder = FlattenMlp(input_size=obs_dim,
output_size=latent_dim,
**encoder_kwargs)
def make_qf():
if have_no_disentangled_encoder:
return FlattenMlp(
input_size=obs_dim + goal_dim + action_dim,
output_size=1,
**qf_kwargs,
)
else:
return DisentangledMlpQf(goal_processor=encoder,
preprocess_obs_dim=obs_dim,
action_dim=action_dim,
qf_kwargs=qf_kwargs,
vectorized=vectorized,
**disentangled_qf_kwargs)
qf1 = make_qf()
qf2 = make_qf()
target_qf1 = make_qf()
target_qf2 = make_qf()
policy = TanhGaussianPolicy(obs_dim=obs_dim + goal_dim,
action_dim=action_dim,
**policy_kwargs)
replay_buffer = ObsDictRelabelingBuffer(
env=train_env,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
achieved_goal_key=achieved_goal_key,
vectorized=vectorized,
**replay_buffer_kwargs)
sac_trainer = SACTrainer(env=train_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**twin_sac_trainer_kwargs)
trainer = HERTrainer(sac_trainer)
eval_path_collector = VAEWrappedEnvPathCollector(
eval_env,
MakeDeterministic(policy),
max_path_length,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
goal_sampling_mode=vae_evaluation_goal_sampling_mode,
)
expl_path_collector = VAEWrappedEnvPathCollector(
train_env,
policy,
max_path_length,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
goal_sampling_mode=vae_exploration_goal_sampling_mode,
)
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=train_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
max_path_length=max_path_length,
**algo_kwargs,
)
algorithm.to(ptu.device)
if save_video:
save_vf_heatmap = save_video_kwargs.get('save_vf_heatmap', True)
if have_no_disentangled_encoder:
def v_function(obs):
action = policy.get_actions(obs)
obs, action = ptu.from_numpy(obs), ptu.from_numpy(action)
return qf1(obs, action)
add_heatmap = partial(add_heatmap_img_to_o_dict,
v_function=v_function)
else:
def v_function(obs):
action = policy.get_actions(obs)
obs, action = ptu.from_numpy(obs), ptu.from_numpy(action)
return qf1(obs, action, return_individual_q_vals=True)
add_heatmap = partial(
add_heatmap_imgs_to_o_dict,
v_function=v_function,
vectorized=vectorized,
)
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,
full_o_postprocess_func=add_heatmap if save_vf_heatmap else None,
)
img_keys = ['v_vals'] + [
'v_vals_dim_{}'.format(dim) for dim in range(latent_dim)
]
eval_video_func = get_save_video_function(rollout_function,
eval_env,
MakeDeterministic(policy),
get_extra_imgs=partial(
get_extra_imgs,
img_keys=img_keys),
tag="eval",
**save_video_kwargs)
train_video_func = get_save_video_function(rollout_function,
train_env,
policy,
get_extra_imgs=partial(
get_extra_imgs,
img_keys=img_keys),
tag="train",
**save_video_kwargs)
algorithm.post_train_funcs.append(eval_video_func)
algorithm.post_train_funcs.append(train_video_func)
algorithm.train()
def experiment(variant):
expl_env = FlatEnv(PointmassBaseEnv(observation_mode='pixels',
transpose_image=True),
use_robot_state=False)
eval_env = expl_env
img_width, img_height = (48, 48)
num_channels = 3
action_dim = int(np.prod(eval_env.action_space.shape))
cnn_params = variant['cnn_params']
cnn_params.update(
input_width=img_width,
input_height=img_height,
input_channels=num_channels,
added_fc_input_size=4,
output_conv_channels=True,
output_size=None,
)
qf_cnn = CNN(**cnn_params)
qf_obs_processor = nn.Sequential(
qf_cnn,
Flatten(),
)
qf_kwargs = copy.deepcopy(variant['qf_kwargs'])
qf_kwargs['obs_processor'] = qf_obs_processor
qf_kwargs['output_size'] = 1
qf_kwargs['input_size'] = (action_dim + qf_cnn.conv_output_flat_size)
qf1 = MlpQfWithObsProcessor(**qf_kwargs)
qf2 = MlpQfWithObsProcessor(**qf_kwargs)
target_qf_cnn = CNN(**cnn_params)
target_qf_obs_processor = nn.Sequential(
target_qf_cnn,
Flatten(),
)
target_qf_kwargs = copy.deepcopy(variant['qf_kwargs'])
target_qf_kwargs['obs_processor'] = target_qf_obs_processor
target_qf_kwargs['output_size'] = 1
target_qf_kwargs['input_size'] = (action_dim +
target_qf_cnn.conv_output_flat_size)
target_qf1 = MlpQfWithObsProcessor(**target_qf_kwargs)
target_qf2 = MlpQfWithObsProcessor(**target_qf_kwargs)
action_dim = int(np.prod(eval_env.action_space.shape))
policy_cnn = CNN(**cnn_params)
policy_obs_processor = nn.Sequential(
policy_cnn,
Flatten(),
)
policy = TanhGaussianPolicyAdapter(policy_obs_processor,
policy_cnn.conv_output_flat_size,
action_dim, **variant['policy_kwargs'])
eval_policy = MakeDeterministic(policy)
eval_path_collector = MdpPathCollector(
eval_env, eval_policy, **variant['eval_path_collector_kwargs'])
replay_buffer = EnvReplayBuffer(
variant['replay_buffer_size'],
expl_env,
)
trainer = SACTrainer(env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**variant['trainer_kwargs'])
if variant['collection_mode'] == 'batch':
expl_path_collector = MdpPathCollector(
expl_env, policy, **variant['expl_path_collector_kwargs'])
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'])
elif variant['collection_mode'] == 'online':
expl_path_collector = MdpStepCollector(
expl_env, policy, **variant['expl_path_collector_kwargs'])
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,
**variant['algo_kwargs'])
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
eval_env = gym.make('FetchReach-v1')
expl_env = gym.make('FetchReach-v1')
observation_key = 'observation'
desired_goal_key = 'desired_goal'
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 = FlattenMlp(input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs'])
qf2 = FlattenMlp(input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs'])
target_qf1 = FlattenMlp(input_size=obs_dim + action_dim + goal_dim,
output_size=1,
**variant['qf_kwargs'])
target_qf2 = FlattenMlp(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'])
eval_policy = MakeDeterministic(policy)
trainer = SACTrainer(env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**variant['sac_trainer_kwargs'])
trainer = HERTrainer(trainer)
eval_path_collector = GoalConditionedPathCollector(
eval_env,
eval_policy,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
)
expl_path_collector = GoalConditionedPathCollector(
expl_env,
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 _use_disentangled_encoder_distance(
max_path_length,
encoder_kwargs,
disentangled_qf_kwargs,
qf_kwargs,
sac_trainer_kwargs,
replay_buffer_kwargs,
policy_kwargs,
evaluation_goal_sampling_mode,
exploration_goal_sampling_mode,
algo_kwargs,
env_id=None,
env_class=None,
env_kwargs=None,
encoder_key_prefix='encoder',
encoder_input_prefix='state',
latent_dim=2,
reward_mode=EncoderWrappedEnv.ENCODER_DISTANCE_REWARD,
# Video parameters
save_video=True,
save_video_kwargs=None,
save_vf_heatmap=True,
**kwargs
):
if save_video_kwargs is None:
save_video_kwargs = {}
if env_kwargs is None:
env_kwargs = {}
assert env_id or env_class
vectorized = (
reward_mode == EncoderWrappedEnv.VECTORIZED_ENCODER_DISTANCE_REWARD
)
if env_id:
import gym
import multiworld
multiworld.register_all_envs()
raw_train_env = gym.make(env_id)
raw_eval_env = gym.make(env_id)
else:
raw_eval_env = env_class(**env_kwargs)
raw_train_env = env_class(**env_kwargs)
raw_train_env.goal_sampling_mode = exploration_goal_sampling_mode
raw_eval_env.goal_sampling_mode = evaluation_goal_sampling_mode
raw_obs_dim = (
raw_train_env.observation_space.spaces['state_observation'].low.size
)
action_dim = raw_train_env.action_space.low.size
encoder = FlattenMlp(
input_size=raw_obs_dim,
output_size=latent_dim,
**encoder_kwargs
)
encoder = Identity()
encoder.input_size = raw_obs_dim
encoder.output_size = raw_obs_dim
np_encoder = EncoderFromMlp(encoder)
train_env = EncoderWrappedEnv(
raw_train_env, np_encoder, encoder_input_prefix,
key_prefix=encoder_key_prefix,
reward_mode=reward_mode,
)
eval_env = EncoderWrappedEnv(
raw_eval_env, np_encoder, encoder_input_prefix,
key_prefix=encoder_key_prefix,
reward_mode=reward_mode,
)
observation_key = '{}_observation'.format(encoder_key_prefix)
desired_goal_key = '{}_desired_goal'.format(encoder_key_prefix)
achieved_goal_key = '{}_achieved_goal'.format(encoder_key_prefix)
obs_dim = train_env.observation_space.spaces[observation_key].low.size
goal_dim = train_env.observation_space.spaces[desired_goal_key].low.size
def make_qf():
return DisentangledMlpQf(
goal_processor=encoder,
preprocess_obs_dim=obs_dim,
action_dim=action_dim,
qf_kwargs=qf_kwargs,
vectorized=vectorized,
**disentangled_qf_kwargs
)
qf1 = make_qf()
qf2 = make_qf()
target_qf1 = make_qf()
target_qf2 = make_qf()
policy = TanhGaussianPolicy(
obs_dim=obs_dim + goal_dim,
action_dim=action_dim,
**policy_kwargs
)
replay_buffer = ObsDictRelabelingBuffer(
env=train_env,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
achieved_goal_key=achieved_goal_key,
vectorized=vectorized,
**replay_buffer_kwargs
)
sac_trainer = SACTrainer(
env=train_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**sac_trainer_kwargs
)
trainer = HERTrainer(sac_trainer)
eval_path_collector = GoalConditionedPathCollector(
eval_env,
MakeDeterministic(policy),
max_path_length,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
goal_sampling_mode='env',
)
expl_path_collector = GoalConditionedPathCollector(
train_env,
policy,
max_path_length,
observation_key=observation_key,
desired_goal_key=desired_goal_key,
goal_sampling_mode='env',
)
algorithm = TorchBatchRLAlgorithm(
trainer=trainer,
exploration_env=train_env,
evaluation_env=eval_env,
exploration_data_collector=expl_path_collector,
evaluation_data_collector=eval_path_collector,
replay_buffer=replay_buffer,
max_path_length=max_path_length,
**algo_kwargs
)
algorithm.to(ptu.device)
if save_video:
def v_function(obs):
action = policy.get_actions(obs)
obs, action = ptu.from_numpy(obs), ptu.from_numpy(action)
return qf1(obs, action, return_individual_q_vals=True)
add_heatmap = partial(
add_heatmap_imgs_to_o_dict,
v_function=v_function,
vectorized=vectorized,
)
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,
full_o_postprocess_func=add_heatmap if save_vf_heatmap else None,
)
img_keys = ['v_vals'] + [
'v_vals_dim_{}'.format(dim) for dim
in range(latent_dim)
]
eval_video_func = get_save_video_function(
rollout_function,
eval_env,
MakeDeterministic(policy),
get_extra_imgs=partial(get_extra_imgs, img_keys=img_keys),
tag="eval",
**save_video_kwargs
)
train_video_func = get_save_video_function(
rollout_function,
train_env,
policy,
get_extra_imgs=partial(get_extra_imgs, img_keys=img_keys),
tag="train",
**save_video_kwargs
)
algorithm.post_train_funcs.append(eval_video_func)
algorithm.post_train_funcs.append(train_video_func)
algorithm.train()
def experiment(variant):
env_params = ENV_PARAMS[variant['env']]
variant.update(env_params)
expl_env = NormalizedBoxEnv(variant['env_class']())
eval_env = NormalizedBoxEnv(variant['env_class']())
obs_dim = expl_env.observation_space.low.size
action_dim = eval_env.action_space.low.size
M = variant['layer_size']
qf1 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[M, M],
)
qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[M, M],
)
target_qf1 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[M, M],
)
target_qf2 = FlattenMlp(
input_size=obs_dim + action_dim,
output_size=1,
hidden_sizes=[M, M],
)
policy = TanhGaussianPolicy(
obs_dim=obs_dim,
action_dim=action_dim,
hidden_sizes=[M, M],
)
eval_policy = MakeDeterministic(policy)
eval_path_collector = MdpPathCollector(
eval_env,
eval_policy,
)
replay_buffer = EnvReplayBuffer(
variant['replay_buffer_size'],
expl_env,
)
trainer = SACTrainer(
env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**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:
expl_path_collector = MdpPathCollector(
expl_env,
policy,
)
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)
algorithm.train()
def experiment(variant):
expl_env = roboverse.make(variant['env'],
gui=False,
randomize=True,
observation_mode='state',
reward_type='shaped',
transpose_image=True)
eval_env = expl_env
action_dim = int(np.prod(eval_env.action_space.shape))
state_dim = eval_env.observation_space.shape[0]
qf_kwargs = copy.deepcopy(variant['qf_kwargs'])
qf_kwargs['output_size'] = 1
qf_kwargs['input_size'] = action_dim + state_dim
qf1 = MlpQf(**qf_kwargs)
qf2 = MlpQf(**qf_kwargs)
target_qf_kwargs = copy.deepcopy(qf_kwargs)
target_qf1 = MlpQf(**target_qf_kwargs)
target_qf2 = MlpQf(**target_qf_kwargs)
policy_kwargs = copy.deepcopy(variant['policy_kwargs'])
policy_kwargs['action_dim'] = action_dim
policy_kwargs['obs_dim'] = state_dim
policy = TanhGaussianPolicy(**policy_kwargs)
eval_policy = MakeDeterministic(policy)
eval_path_collector = MdpPathCollector(
eval_env, eval_policy, **variant['eval_path_collector_kwargs'])
replay_buffer = EnvReplayBuffer(
variant['replay_buffer_size'],
expl_env,
)
trainer = SACTrainer(env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
**variant['trainer_kwargs'])
if variant['scripted_policy']:
if 'V3-v0' in variant['env']:
scripted_policy = GraspV3ScriptedPolicy(
expl_env, noise_std=variant['scripted_noise_std'])
elif 'V4-v0' in variant['env']:
scripted_policy = GraspV4ScriptedPolicy(
expl_env, noise_std=variant['scripted_noise_std'])
elif 'V5-v0' in variant['env']:
scripted_policy = GraspV5ScriptedPolicy(
expl_env, noise_std=variant['scripted_noise_std'])
else:
raise NotImplementedError
else:
scripted_policy = None
if variant['collection_mode'] == 'batch':
expl_path_collector = MdpPathCollector(
expl_env,
policy,
optional_expl_policy=scripted_policy,
optional_expl_probability_init=0.5,
**variant['expl_path_collector_kwargs'])
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'])
elif variant['collection_mode'] == 'online':
expl_path_collector = MdpStepCollector(
expl_env,
policy,
optional_expl_policy=scripted_policy,
optional_expl_probability=0.9,
**variant['expl_path_collector_kwargs'])
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,
**variant['algo_kwargs'])
dump_buffer_func = BufferSaveFunction(variant)
# algorithm.post_train_funcs.append(dump_buffer_func)
algorithm.to(ptu.device)
algorithm.train()