//  [python]
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