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', 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)
    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', 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)
    algorithm.to(ptu.device)
    algorithm.train()