code: def testPredictedRewards(self):
      policy = greedy_multi_objective_policy.GreedyMultiObjectiveNeuralPolicy(
          self._time_step_spec,
          self._action_spec,
          self._scalarizer,
          self._create_objective_networks(),
          emit_policy_info=(
              policy_utilities.InfoFields.PREDICTED_REWARDS_MEAN,
              policy_utilities.InfoFields.MULTIOBJECTIVE_SCALARIZED_PREDICTED_REWARDS_MEAN,
          ),
      )
      observations = tf.constant([[1, 2], [2, 1]], dtype=tf.float32)
      time_step = ts.restart(observations, batch_size=2)
      action_step = policy.action(time_step)
      self.assertEqual(action_step.action.shape.as_list(), [2])
      self.assertEqual(action_step.action.dtype, tf.int32)
      # Initialize all variables
      self.evaluate(tf.compat.v1.global_variables_initializer())
      self.assertAllEqual(self.evaluate(action_step.action), [2, 0])
      # The expected values are obtained by passing the observation through the
      # Keras dense layer of the DummyNet (defined above).
      predicted_rewards_expected_array = np.array(
          [
              [[8, 11, 14], [12, 8, 13], [11, 14, 8]],
              [[5, 8, 11], [10, 5, 9], [8, 11, 5]],
          ]
      )
      p_info = self.evaluate(action_step.info)
      predicted_rewards = getattr(
          p_info, policy_utilities.InfoFields.PREDICTED_REWARDS_MEAN
      )
      self.assertAllClose(
          predicted_rewards, predicted_rewards_expected_array
      )
      self.assertAllClose(
          getattr(
              p_info,
              policy_utilities.InfoFields.MULTIOBJECTIVE_SCALARIZED_PREDICTED_REWARDS_MEAN,
          ),
          greedy_multi_objective_policy.scalarize_objectives(
              predicted_rewards, policy.scalarizer
          ),
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  from typing import Dict, List, Text, Tuple
  import numpy
  import tensorflow
  from tf_agents.bandits.multi_objective import multi_objective_scalarizer
  from tf_agents.bandits.networks import global_and_arm_feature_network
  from tf_agents.bandits.networks import heteroscedastic_q_network
  from tf_agents.bandits.policies import greedy_multi_objective_neural_policy
  from tf_agents.bandits.specs import utils
  from tf_agents.networks import network
  from tf_agents.policies import utils
  from tf_agents.specs import array_spec
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import test_utils
label_desc: Use image node_id 17 for calling the GreedyRewardPredictionPolicyTest obj's underlying member method code with example usage: obj.testPredictedRewards() without return types