def experiment(variant):
#expl_env = carla_env.CarlaObsDictEnv(args=variant['env_args'])
import gym
import d4rl.carla
expl_env = gym.make('carla-lane-dict-v0')
eval_env = expl_env
#num_channels, img_width, img_height = eval_env._wrapped_env.image_shape
num_channels, img_width, img_height = eval_env.image_shape
# num_channels = 3
action_dim = int(np.prod(eval_env.action_space.shape))
# obs_dim = 11
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'])
eval_policy = MakeDeterministic(policy)
observation_key = 'image'
eval_path_collector = ObsDictPathCollector(
eval_env,
eval_policy,
observation_key=observation_key,
**variant['eval_path_collector_kwargs'])
expl_path_collector = CustomObsDictPathCollector(
expl_env,
observation_key=observation_key,
)
observation_key = 'image'
replay_buffer = ObsDictReplayBuffer(
variant['replay_buffer_size'],
expl_env,
observation_key=observation_key,
)
load_hdf5(expl_env, replay_buffer)
#load_buffer(buffer_path=variant['buffer'], replay_buffer=replay_buffer)
# import ipdb; ipdb.set_trace()
trainer = SACTrainer(env=eval_env,
policy=policy,
qf1=qf1,
qf2=qf2,
target_qf1=target_qf1,
target_qf2=target_qf2,
behavior_policy=None,
**variant['trainer_kwargs'])
variant['algo_kwargs']['max_path_length'] = expl_env._max_episode_steps
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,
eval_both=True,
batch_rl=True,
**variant['algorithm_kwargs'])
video_func = VideoSaveFunctionBullet(variant)
algorithm.post_train_funcs.append(video_func)
algorithm.to(ptu.device)
algorithm.train()
def experiment(variant):
expl_env = gym.make('carla-lane-dict-v0')
eval_env = expl_env
num_channels, 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']
)
cnn_vae_params = variant['cnn_vae_params']
cnn_vae_params['conv_args'].update(
input_width=img_width,
input_height=img_height,
input_channels=num_channels,
)
vae_policy = ConvVAEPolicy(
representation_size=cnn_vae_params['representation_size'],
architecture=cnn_vae_params,
action_dim=action_dim,
input_channels=3,
imsize=img_width,
)
observation_key = 'image'
eval_path_collector = CustomObsDictPathCollector(
eval_env,
observation_key=observation_key,
**variant['eval_path_collector_kwargs']
)
vae_eval_path_collector = CustomObsDictPathCollector(
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)
observation_key = 'image'
replay_buffer = ObsDictReplayBuffer(
variant['replay_buffer_size'],
expl_env,
observation_key=observation_key,
)
load_hdf5(expl_env, replay_buffer)
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 = 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,
vae_evaluation_data_collector=vae_eval_path_collector,
replay_buffer=replay_buffer,
q_learning_alg=True,
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=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))
# obs_dim = 11
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'])
eval_policy = MakeDeterministic(policy)
observation_key = 'image'
eval_path_collector = ObsDictPathCollector(
eval_env,
eval_policy,
observation_key=observation_key,
**variant['eval_path_collector_kwargs'])
expl_path_collector = CustomObsDictPathCollector(
expl_env,
observation_key=observation_key,
)
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,
behavior_policy=None,
**variant['trainer_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,
eval_both=True,
batch_rl=variant['load_buffer'],
**variant['algorithm_kwargs'])
video_func = VideoSaveFunctionBullet(variant)
algorithm.post_train_funcs.append(video_func)
algorithm.to(ptu.device)
algorithm.train()