code: def testPerArmRewards(self):
      policy, observations = self._create_arm_policy_and_observations()
      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())
      action, p_info, first_arm_features = self.evaluate(
          [
              action_step.action,
              action_step.info,
              observations[bandit_spec_utils.PER_ARM_FEATURE_KEY][0],
          ]
      )
      self.assertAllEqual(action.shape, [2])
      predicted_rewards = getattr(
          p_info, policy_utilities.InfoFields.PREDICTED_REWARDS_MEAN
      )
      self.assertAllEqual(predicted_rewards.shape, [2, 3, 4])
      scalarized_predicted_rewards = getattr(
          p_info,
          policy_utilities.InfoFields.MULTIOBJECTIVE_SCALARIZED_PREDICTED_REWARDS_MEAN,
      )
      self.assertAllEqual(scalarized_predicted_rewards.shape, [2, 4])
      chosen_arm_features = getattr(
          p_info, policy_utilities.InfoFields.CHOSEN_ARM_FEATURES
      )
      self.assertAllEqual(chosen_arm_features.shape, [2, 3])
      first_action = action[0]
      self.assertAllEqual(
          chosen_arm_features[0], first_arm_features[first_action]
      )
  
      # Check that zeroing out some of the actions does not affect the predicted
      # rewards for unchanged actions. This is to make sure that action feature
      # padding does not influence the behavior.
  
      if not tf.executing_eagerly():
          # The below comparison will only work in tf2 because of the random per-arm
          # observations get re-drawn in tf1.
          return
      action_feature = observations[
          bandit_spec_utils.PER_ARM_FEATURE_KEY
      ]
      padded_action_feature = tf.concat(
          [
              action_feature[:, 0:1, :],
              tf.zeros(shape=[2, 3, 3], dtype=tf.float32),
          ],
          axis=1,
      )
      observations[
          bandit_spec_utils.PER_ARM_FEATURE_KEY
      ] = padded_action_feature
      time_step = ts.restart(observations, batch_size=2)
      padded_action_step = policy.action(time_step)
      padded_p_info = self.evaluate(padded_action_step.info)
      self.assertAllEqual(
          predicted_rewards[:, :, 0],
          getattr(
              padded_p_info,
              policy_utilities.InfoFields.PREDICTED_REWARDS_MEAN,
          )[:, :, 0],
      )
      self.assertAllEqual(
          scalarized_predicted_rewards[:, 0],
          getattr(
              padded_p_info,
              policy_utilities.InfoFields.MULTIOBJECTIVE_SCALARIZED_PREDICTED_REWARDS_MEAN,
          )[:, 0],
      )
  
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 19 for calling the GreedyRewardPredictionPolicyTest obj's underlying member method code with example usage: obj.testPerArmRewards() without return types