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: Distributed train fn to be passed as input to run().
code: def _distributed_train_step(self, experience, weights=None):
      """Distributed train fn to be passed as input to run()."""
      (
          experience_reward,
          action,
          experience_observation,
          batch_size,
      ) = self._process_experience(experience)
      self._train_step_counter.assign_add(batch_size)
      observation, reward = self._maybe_apply_per_example_weight(
          experience_observation, experience_reward, weights
      )
      loss = self._loss(experience, weights, training=True)
  
      for k in range(self._num_models):
          diag_mask = tf.linalg.tensor_diag(
              tf.cast(tf.equal(action, k), self._dtype)
          )
          observations_for_arm = tf.matmul(diag_mask, observation)
          rewards_for_arm = tf.matmul(
              diag_mask, tf.reshape(reward, [-1, 1])
          )
  
          # Compute local updates for the matrix A and b of this arm.
          cov_matrix_local_update = tf.matmul(
              observations_for_arm,
              observations_for_arm,
              transpose_a=True,
          )
          data_vector_local_update = (
              bandit_utils.sum_reward_weighted_observations(
                  rewards_for_arm, observations_for_arm
              )
          )
  
          def _merge_fn(
              strategy,
              per_replica_cov_matrix_update,
              per_replica_data_vector_update,
          ):
              """Merge the per-replica-updates."""
              # Reduce the per-replica-updates using SUM.
              # pylint: disable=cell-var-from-loop
              updates_and_vars = [
                  (
                      per_replica_cov_matrix_update,
                      self._cov_matrix_list[k],
                  ),
                  (
                      per_replica_data_vector_update,
                      self._data_vector_list[k],
                  ),
              ]
  
              reduced_updates = strategy.extended.batch_reduce_to(
                  tf.distribute.ReduceOp.SUM, updates_and_vars
              )
  
              # Update the model variables.
              self._cov_matrix_list[k].assign_add(reduced_updates[0])
              self._data_vector_list[k].assign_add(reduced_updates[1])
  
              # Compute the eigendecomposition, if needed.
              if self._use_eigendecomp:
                  eig_vals, eig_matrix = tf.linalg.eigh(
                      self._cov_matrix_list[k]
                  )
                  self._eig_vals_list[k].assign(eig_vals)
                  self._eig_matrix_list[k].assign(eig_matrix)
  
          # Passes the local_updates to the _merge_fn() above that performs custom
          # computation on the per-replica values.
          # All replicas pause their execution until merge_call() is done and then,
          # execution is resumed.
          replica_context = tf.distribute.get_replica_context()
          replica_context.merge_call(
              _merge_fn,
              args=(cov_matrix_local_update, data_vector_local_update),
          )
  
      with tf.control_dependencies([loss.loss]):
          self._theta.assign(self.theta)
      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 18 for calling the LinearBanditAgent obj's underlying member method code with example usage: obj._distributed_train_step(experience, weights) and returns: loss