class_docstr: Manages all the learning details needed when training an agent.
  
  These include:
    * Using distribution strategies correctly
    * Summaries
    * Checkpoints
    * Minimizing entering/exiting TF context:
        Especially in the case of TPUs scheduling a single TPU program to
        perform multiple train steps is critical for performance.
    * Generalizes the train call to be done correctly across CPU, GPU, or TPU
      executions managed by DistributionStrategies. This uses `strategy.run` and
      then makes sure to do a reduce operation over the `LossInfo` returned by
      the agent.
class_docstr_tok: ["Manages","all","the","learning","details","needed","when","training","an","agent",".","These","include",":","*","Using","distribution","strategies","correctly","*","Summaries","*","Checkpoints","*","Minimizing","entering\/exiting","TF","context",":","Especially","in","the","case","of","TPUs","scheduling","a","single","TPU","program","to","perform","multiple","train","steps","is","critical","for","performance",".","*","Generalizes","the","train","call","to","be","done","correctly","across","CPU",",","GPU",",","or","TPU","executions","managed","by","DistributionStrategies",".","This","uses","`","strategy.run","`","and","then","makes","sure","to","do","a","reduce","operation","over","the","`","LossInfo","`","returned","by","the","agent","."]
code: def _loss(
      self,
      experience_and_sample_info: ExperienceAndSampleInfo,
      reduce_op: tf.distribute.ReduceOp,
  ) -> tf_agent.LossInfo:
      (experience, sample_info) = experience_and_sample_info
  
      if self.use_kwargs_in_agent_train:
          loss_info = self.strategy.run(
              self._agent.loss, kwargs=experience
          )
      else:
          loss_info = self.strategy.run(
              self._agent.loss, args=(experience,)
          )
  
      if self.after_train_strategy_step_fn:
          if self.use_kwargs_in_agent_train:
              self.strategy.run(
                  self.after_train_strategy_step_fn,
                  kwargs=dict(
                      experience=(experience, sample_info),
                      loss_info=loss_info,
                  ),
                  options=self._strategy_run_options,
              )
          else:
              self.strategy.run(
                  self.after_train_strategy_step_fn,
                  args=((experience, sample_info), loss_info),
                  options=self._strategy_run_options,
              )
  
      def _reduce_loss(loss):
          rank = None
          if isinstance(loss, tf.distribute.DistributedValues):
              rank = loss.values[0].shape.rank
          elif tf.is_tensor(loss):
              rank = loss.shape.rank
          axis = None
          if rank:
              axis = tuple(range(0, rank))
          return self.strategy.reduce(reduce_op, loss, axis=axis)
  
      # We assume all data can be reduced in the loss_info. This means no
      # string dtypes are currently allowed as LossInfo Fields.
      reduced_loss_info = tf.nest.map_structure(_reduce_loss, loss_info)
      return reduced_loss_info
  
file_path: agents/tf_agents/train/learner.py
imports: import os
  from typing import Any, Optional, Tuple
  from absl import logging
  import gin
  import tensorflow.compat.v2
  from tf_agents.agents import tf_agent
  from tf_agents.specs import tensor_spec
  from tf_agents.train import interval_trigger
  from tf_agents.typing import types
  from tf_agents.utils import common
label_desc: Use image node_id 8 for calling the Learner obj's underlying member method code with example usage: obj._loss(experience_and_sample_info, reduce_op) and returns: reduced_loss_info, self