code: def get_environment_and_optimal_functions_by_name(
      environment_name, batch_size
  ):
      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)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import os
  import tempfile
  from unittest import mock
  from absl import logging
  from absl.testing import parameterized
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.agents import exp3_agent
  from tf_agents.bandits.agents.examples.v2 import trainer
  from tf_agents.bandits.agents.examples.v2 import trainer_test_utils
  from tf_agents.bandits.environments import environment_utilities
  from tf_agents.bandits.environments import random_bandit_environment
  from tf_agents.bandits.environments import stationary_stochastic_py_environment
  from tf_agents.bandits.environments import wheel_py_environment
  from tf_agents.bandits.metrics import tf_metrics
  from tf_agents.environments import tf_py_environment
  from tf_agents.metrics import export_utils
  from tf_agents.specs import tensor_spec
label_desc: Use image node_id 2 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