code: def compute_optimal_reward(
      observation, per_action_reward_fns, enable_noise=False
  ):
      """Computes the optimal reward.
  
      Args:
        observation: a (possibly batched) observation.
        per_action_reward_fns: a list of reward functions; one per action. Each
          reward function generates a reward when called with an observation.
        enable_noise: (bool) whether to add noise to the rewards.
  
      Returns:
        The optimal reward.
      """
      num_actions = len(per_action_reward_fns)
      rewards = np.stack(
          [
              per_action_reward_fns[a](observation, enable_noise)
              for a in range(num_actions)
          ],
          axis=-1,
      )
      # `rewards` should be of shape [`batch_size`, `num_actions`].
      optimal_action_reward = np.max(rewards, axis=-1)
      return optimal_action_reward
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import gin
  import numpy
  from six.moves import range
  import tensorflow
  from tf_agents.bandits.agents import utils
  from tf_agents.bandits.environments import wheel_py_environment
label_desc: Use image node_id 9 for calling a global function with example usage: compute_optimal_reward(observation, per_action_reward_fns, enable_noise) and returns: optimal_action_reward