code: def testPredictedRewards(
      self, batch_size, exploration_strategy, alpha
  ):
      a_inv = []
      theta = []
      for k in range(self._num_actions):
          a_inv.append(
              np.linalg.inv(self._a_numpy[k] + np.eye(self._obs_dim))
          )
          theta.append(
              np.matmul(
                  a_inv[k], self._b_numpy[k].reshape([self._obs_dim, 1])
              )
          )
      policy = linear_policy.LinearBanditPolicy(
          self._action_spec,
          self._a,
          self._b,
          self._num_samples_per_arm,
          self._time_step_spec,
          exploration_strategy,
          alpha=alpha,
          theta=None
          if alpha > 0.0
          else tf.squeeze(
              tf.cast(tf.stack(theta), dtype=self._obs_spec.dtype),
              axis=-1,
          ),
          emit_policy_info=(
              policy_utilities.InfoFields.PREDICTED_REWARDS_MEAN,
          ),
      )
  
      action_step = policy.action(
          self._time_step_batch(batch_size=batch_size)
      )
      self.assertEqual(action_step.action.shape.as_list(), [batch_size])
      self.assertEqual(action_step.action.dtype, tf.int32)
  
      observation_numpy = np.array(
          range(batch_size * self._obs_dim), dtype=np.float32
      ).reshape([batch_size, self._obs_dim])
  
      predicted_rewards_expected = []
      for k in range(self._num_actions):
          est_mean_reward = np.matmul(observation_numpy, theta[k])
          predicted_rewards_expected.append(est_mean_reward)
  
      predicted_rewards_expected_array = np.stack(
          predicted_rewards_expected, axis=-1
      ).reshape(batch_size, self._num_actions)
      p_info = self.evaluate(action_step.info)
      self.assertAllClose(
          p_info.predicted_rewards_mean,
          predicted_rewards_expected_array,
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import itertools
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents.bandits.policies import linear_bandit_policy
  from tf_agents.bandits.specs import utils
  from tf_agents.policies import utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.utils import test_utils
label_desc: Use image node_id 19 for calling the LinearBanditPolicyTest obj's underlying member method code with example usage: obj.testPredictedRewards(batch_size, exploration_strategy, alpha) without return types