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 = AWRSACTrainer(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):
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 experiment(variant):
expl_env = roboverse.make(variant['env'],
gui=False,
randomize=True,
observation_mode='state',
reward_type='shaped')
eval_env = expl_env
action_dim = int(np.prod(eval_env.action_space.shape))
state_dim = obs_dim = 11
M = 256
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)
vae_policy = VAEPolicy(
obs_dim=obs_dim,
action_dim=action_dim,
hidden_sizes=[M, M],
latent_dim=action_dim * 2,
)
eval_path_collector = CustomMdpPathCollector(
eval_env, save_images=True, **variant['eval_path_collector_kwargs'])
vae_eval_path_collector = CustomMdpPathCollector(
eval_env,
max_num_epoch_paths_saved=5,
save_images=True,
)
with open(variant['buffer'], 'rb') as f:
replay_buffer = pickle.load(f)
trainer = BEARTrainer(env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
vae=vae_policy,
**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,
vae_evaluation_data_collector=vae_eval_path_collector,
replay_buffer=replay_buffer,
q_learning_alg=True,
batch_rl=variant['batch_rl'],
**variant['algo_kwargs'])
algorithm.to(ptu.device)
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': 0.0,
'arm_object_distance_cost_coeff': 1.0,
'goal': (0, -1),
}
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):
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()