code: def testPerArmObservation(
      self, batch_size, actions_from_reward_layer
  ):
      global_obs_dim = 7
      arm_obs_dim = 3
      obs_spec = bandit_spec_utils.create_per_arm_observation_spec(
          global_obs_dim,
          arm_obs_dim,
          self._num_actions,
          add_num_actions_feature=True,
      )
      time_step_spec = ts.time_step_spec(obs_spec)
      dummy_net = arm_network.create_feed_forward_common_tower_network(
          obs_spec,
          global_layers=(3, 4, 5),
          arm_layers=(3, 2),
          common_layers=(4, 3),
          output_dim=self._encoding_dim,
      )
      reward_layer = get_per_arm_reward_layer(
          encoding_dim=self._encoding_dim
      )
  
      policy = neural_linucb_policy.NeuralLinUCBPolicy(
          dummy_net,
          self._encoding_dim,
          reward_layer,
          actions_from_reward_layer=tf.constant(
              actions_from_reward_layer, dtype=tf.bool
          ),
          cov_matrix=self._a[0:1],
          data_vector=self._b[0:1],
          num_samples=self._num_samples_per_arm[0:1],
          epsilon_greedy=0.0,
          time_step_spec=time_step_spec,
          accepts_per_arm_features=True,
          emit_policy_info=(
              "predicted_rewards_mean",
              "predicted_rewards_optimistic",
          ),
      )
  
      current_time_step = self._per_arm_time_step_batch(
          batch_size=batch_size,
          global_obs_dim=global_obs_dim,
          arm_obs_dim=arm_obs_dim,
      )
      action_step = policy.action(current_time_step)
      self.assertEqual(action_step.action.dtype, tf.int32)
      self.evaluate(tf.compat.v1.global_variables_initializer())
      action_fn = common.function_in_tf1()(policy.action)
      action_step = action_fn(current_time_step)
  
      input_observation = current_time_step.observation
      encoded_observation, _ = dummy_net(input_observation)
  
      if actions_from_reward_layer:
          predicted_rewards_from_reward_layer = reward_layer(
              encoded_observation
          )
          predicted_rewards_expected = self.evaluate(
              predicted_rewards_from_reward_layer
          ).reshape((-1, self._num_actions))
      else:
          observation_numpy = self.evaluate(encoded_observation)
          predicted_rewards_expected = (
              self._get_predicted_rewards_from_per_arm_linucb(
                  observation_numpy, batch_size
              )
          )
  
      p_info = self.evaluate(action_step.info)
      self.assertEqual(p_info.predicted_rewards_mean.dtype, np.float32)
      self.assertAllClose(
          p_info.predicted_rewards_mean, predicted_rewards_expected
      )
      self.assertAllGreaterEqual(
          p_info.predicted_rewards_optimistic
          - predicted_rewards_expected,
          0,
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents.bandits.networks import global_and_arm_feature_network
  from tf_agents.bandits.policies import neural_linucb_policy
  from tf_agents.bandits.specs import utils
  from tf_agents.networks import network
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.utils import common
  from tf_agents.utils import test_utils
label_desc: Use image node_id 18 for calling the NeuralLinUCBPolicyTest obj's underlying member method code with example usage: obj.testPerArmObservation(batch_size, actions_from_reward_layer) without return types