def METHOD_NAME(environment_name, batch_size):
  """Helper function that outputs an environment and related functions.
  Args:
    environment_name: The (string) name of the desired environment.
    batch_size: The batch_size
  Returns:
    A tuple of (environment, optimal_reward_fn, optimal_action_fn), where the
    latter two functions are for calculating regret and the suboptimal actions
    metrics.
  """
  if environment_name == 'stationary_stochastic':
    context_dim = 7
    num_actions = 5
    action_reward_fns = (
        environment_utilities.sliding_linear_reward_fn_generator(
            context_dim, num_actions, 0.1
        )
    )
    py_env = (
        stationary_stochastic_py_environment.StationaryStochasticPyEnvironment(
            functools.partial(
                environment_utilities.context_sampling_fn,
                batch_size=batch_size,
                context_dim=context_dim,
            ),
            action_reward_fns,
            batch_size=batch_size,
        )
    )
    optimal_reward_fn = functools.partial(
        environment_utilities.tf_compute_optimal_reward,
        per_action_reward_fns=action_reward_fns,
    )
    optimal_action_fn = functools.partial(
        environment_utilities.tf_compute_optimal_action,
        per_action_reward_fns=action_reward_fns,
    )
    environment = tf_py_environment.TFPyEnvironment(py_env)
  elif environment_name == 'wheel':
    delta = 0.5
    mu_base = [0.05, 0.01, 0.011, 0.009, 0.012]
    std_base = [0.001] * 5
    mu_high = 0.5
    std_high = 0.001
    py_env = wheel_py_environment.WheelPyEnvironment(
        delta, mu_base, std_base, mu_high, std_high, batch_size
    )
    environment = tf_py_environment.TFPyEnvironment(py_env)
    optimal_reward_fn = functools.partial(
        environment_utilities.tf_wheel_bandit_compute_optimal_reward,
        delta=delta,
        mu_inside=mu_base[0],
        mu_high=mu_high,
    )
    optimal_action_fn = functools.partial(
        environment_utilities.tf_wheel_bandit_compute_optimal_action,
        delta=delta,
    )
  return (environment, optimal_reward_fn, optimal_action_fn)