docstr: Check that the agent updates for specified actions and rewards.
docstr_tok: ["Check","that","the","agent","updates","for","specified","actions","and","rewards","."]
code: def testLinearAgentUpdatePerArmFeatures(
      self,
      batch_size,
      context_dim,
      exploration_policy,
      dtype,
      use_eigendecomp=False,
      set_example_weights=False,
  ):
      """Check that the agent updates for specified actions and rewards."""
  
      # Construct a `Trajectory` for the given action, observation, reward.
      num_actions = 5
      global_context_dim = context_dim
      arm_context_dim = 3
      (
          initial_step,
          final_step,
      ) = _get_initial_and_final_steps_with_per_arm_features(
          batch_size,
          global_context_dim,
          num_actions,
          arm_context_dim,
          num_actions_feature=True,
      )
      action = np.random.randint(
          num_actions, size=batch_size, dtype=np.int32
      )
      action_step = policy_step.PolicyStep(
          action=tf.convert_to_tensor(action),
          info=policy_utilities.PerArmPolicyInfo(
              chosen_arm_features=np.arange(
                  batch_size * arm_context_dim, dtype=np.float32
              ).reshape([batch_size, arm_context_dim])
          ),
      )
      experience = _get_experience(
          initial_step, action_step, final_step
      )
  
      # Construct an agent and perform the update.
      observation_spec = (
          bandit_spec_utils.create_per_arm_observation_spec(
              context_dim,
              arm_context_dim,
              num_actions,
              add_num_actions_feature=True,
          )
      )
      time_step_spec = time_step.time_step_spec(observation_spec)
      action_spec = tensor_spec.BoundedTensorSpec(
          dtype=tf.int32, shape=(), minimum=0, maximum=num_actions - 1
      )
      agent = linear_agent.LinearBanditAgent(
          exploration_policy=exploration_policy,
          time_step_spec=time_step_spec,
          action_spec=action_spec,
          use_eigendecomp=use_eigendecomp,
          accepts_per_arm_features=True,
          dtype=dtype,
      )
      self.evaluate(agent.initialize())
      weights = (
          tf.linspace(start=1.5, stop=10.5, num=batch_size)
          if set_example_weights
          else None
      )
      loss_info = agent.train(experience, weights)
      self.evaluate(loss_info)
      final_a = self.evaluate(agent.cov_matrix)
      final_b = self.evaluate(agent.data_vector)
  
      # Compute the expected updated estimates.
      global_observation = experience.observation[
          bandit_spec_utils.GLOBAL_FEATURE_KEY
      ]
      arm_observation = experience.policy_info.chosen_arm_features
      overall_observation = tf.squeeze(
          tf.concat([global_observation, arm_observation], axis=-1),
          axis=1,
      )
      squeezed_rewards = tf.squeeze(experience.reward, axis=1)
      observation, rewards = _maybe_weight_observation_and_reward(
          overall_observation, squeezed_rewards, weights
      )
      expected_a_new = tf.matmul(
          observation, observation, transpose_a=True
      )
      expected_b_new = bandit_utils.sum_reward_weighted_observations(
          rewards, observation
      )
      self.assertAllClose(expected_a_new, final_a[0])
      self.assertAllClose(expected_b_new, final_b[0])
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  from typing import Optional, Tuple
  from absl import flags
  from absl.testing import parameterized
  import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.agents import lin_ucb_agent
  from tf_agents.bandits.agents import linear_bandit_agent
  from tf_agents.bandits.agents import linear_thompson_sampling_agent
  from tf_agents.bandits.agents import utils
  from tf_agents.bandits.drivers import driver_utils
  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 policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import common
label_desc: Use image node_id 8 for calling the LinearBanditAgentTest obj's underlying member method code with example usage: obj.testLinearAgentUpdatePerArmFeatures(batch_size, context_dim, exploration_policy, dtype, use_eigendecomp, set_example_weights) without return types