code: def get_environment_and_optimal_functions_by_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)
  
file_path: agents/tf_agents/bandits/agents/examples/v2/trainer_test_utils.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.agents import exp3_mixture_agent
  from tf_agents.bandits.agents import lin_ucb_agent
  from tf_agents.bandits.agents import linear_thompson_sampling_agent
  from tf_agents.bandits.agents import neural_epsilon_greedy_agent
  from tf_agents.bandits.environments import environment_utilities
  from tf_agents.bandits.environments import stationary_stochastic_py_environment
  from tf_agents.bandits.environments import wheel_py_environment
  from tf_agents.bandits.networks import global_and_arm_feature_network
  from tf_agents.environments import tf_py_environment
  from tf_agents.networks import q_network
  from tf_agents.policies import utils
label_desc: Use image node_id 1 for calling a global function with example usage: get_environment_and_optimal_functions_by_name(environment_name, batch_size) and returns: environment,  optimal_reward_fn,  optimal_action_fn