class_docstr: An agent that maintains linear reward estimates and their uncertainties.
class_docstr_tok: ["An","agent","that","maintains","linear","reward","estimates","and","their","uncertainties","."]
docstr: Updates the policy based on the data in `experience`.
  
  Note that `experience` should only contain data points that this agent has
  not previously seen. If `experience` comes from a replay buffer, this buffer
  should be cleared between each call to `train`.
  
  Args:
    experience: A batch of experience data in the form of a `Trajectory`.
    weights: Unused.
  
  Returns:
      A `LossInfo` containing the loss *before* the training step is taken.
      In most cases, if `weights` is provided, the entries of this tuple will
      have been calculated with the weights.  Note that each Agent chooses
      its own method of applying weights.
code: def _train(self, experience, weights=None):
      """Updates the policy based on the data in `experience`.
  
      Note that `experience` should only contain data points that this agent has
      not previously seen. If `experience` comes from a replay buffer, this buffer
      should be cleared between each call to `train`.
  
      Args:
        experience: A batch of experience data in the form of a `Trajectory`.
        weights: Unused.
  
      Returns:
          A `LossInfo` containing the loss *before* the training step is taken.
          In most cases, if `weights` is provided, the entries of this tuple will
          have been calculated with the weights.  Note that each Agent chooses
          its own method of applying weights.
      """
      experience = self._as_trajectory(experience)
  
      if tf.distribute.has_strategy():
          return self._distributed_train_step(experience)
  
      loss = self._loss(experience, weights, training=True)
  
      (
          experience_reward,
          action,
          experience_observation,
          batch_size,
      ) = self._process_experience(experience)
      observation, reward = self._maybe_apply_per_example_weight(
          experience_observation, experience_reward, weights
      )
      for k in range(self._num_models):
          action_mask = tf.equal(action, k)
          observations_for_arm = tf.boolean_mask(
              observation, action_mask, axis=0
          )
          rewards_for_arm = tf.boolean_mask(
              tf.reshape(reward, [-1, 1]), action_mask, axis=0
          )
  
          num_samples_for_arm_current = tf.reduce_sum(
              tf.cast(action_mask, self._dtype)
          )
          tf.compat.v1.assign_add(
              self._num_samples_list[k], num_samples_for_arm_current
          )
          num_samples_for_arm_total = self._num_samples_list[
              k
          ].read_value()
  
          # Update the matrix A and b.
          # pylint: disable=cell-var-from-loop,g-long-lambda
          def update(cov_matrix, data_vector):
              return update_a_and_b_with_forgetting(
                  cov_matrix,
                  data_vector,
                  rewards_for_arm,
                  observations_for_arm,
                  self._gamma,
              )
  
          a_new, b_new = tf.cond(
              tf.squeeze(num_samples_for_arm_total) > 0,
              lambda: update(
                  self._cov_matrix_list[k], self._data_vector_list[k]
              ),
              lambda: (
                  self._cov_matrix_list[k],
                  self._data_vector_list[k],
              ),
          )
          tf.compat.v1.assign(self._cov_matrix_list[k], a_new)
          tf.compat.v1.assign(self._data_vector_list[k], b_new)
  
      if self._use_eigendecomp:
          eigenvalues, eigenvectors = tf.linalg.eigh(
              tf.stack(self._cov_matrix_list)
          )
          eigenvalues_list = tf.unstack(eigenvalues)
          eigenvectors_list = tf.unstack(eigenvectors)
          for k in range(self._num_models):
              tf.compat.v1.assign(
                  self._eig_vals_list[k], eigenvalues_list[k]
              )
              tf.compat.v1.assign(
                  self._eig_matrix_list[k], eigenvectors_list[k]
              )
      else:
          for k in range(self._num_models):
              tf.compat.v1.assign(
                  self._eig_vals_list[k],
                  tf.constant([], dtype=self._dtype),
              )
              tf.compat.v1.assign(
                  self._eig_matrix_list[k],
                  tf.constant([], dtype=self._dtype),
              )
  
      with tf.control_dependencies([loss.loss]):
          self._theta.assign(self.theta)
      self._train_step_counter.assign_add(batch_size)
      return loss
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import enum
  from typing import Optional, Sequence, Text, Tuple
  import gin
  import tensorflow
  from tf_agents.agents import data_converter
  from tf_agents.agents import tf_agent
  from tf_agents.bandits.agents import utils
  from tf_agents.bandits.policies import linalg
  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.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 19 for calling the LinearBanditAgent obj's underlying member method code with example usage: obj._train(experience, weights) and returns: loss, self, update_a_and_b_with_forgetting