code: def tf_compute_optimal_reward(
      observation, per_action_reward_fns, enable_noise=False
  ):
      """TF wrapper around `compute_optimal_reward` to be used in `tf_metrics`."""
      compute_optimal_reward_fn = functools.partial(
          compute_optimal_reward,
          per_action_reward_fns=per_action_reward_fns,
          enable_noise=enable_noise,
      )
      return tf.py_function(
          compute_optimal_reward_fn, [observation], tf.float32
      )
  
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 10 for calling a global function with example usage: tf_compute_optimal_reward(observation, per_action_reward_fns, enable_noise) and returns: tf