class_docstr: A DQN Agent.
  
  Implements the DQN algorithm from
  
  "Human level control through deep reinforcement learning"
    Mnih et al., 2015
    https://deepmind.com/research/dqn/
  
  This agent also implements n-step updates. See "Rainbow: Combining
  Improvements in Deep Reinforcement Learning" by Hessel et al., 2017, for a
  discussion on its benefits: https://arxiv.org/abs/1710.02298
class_docstr_tok: ["A","DQN","Agent",".","Implements","the","DQN","algorithm","from","''","Human","level","control","through","deep","reinforcement","learning","''","Mnih","et","al.",",","2015","https",":","\/\/deepmind.com\/research\/dqn\/","This","agent","also","implements","n-step","updates",".","See","``","Rainbow",":","Combining","Improvements","in","Deep","Reinforcement","Learning","''","by","Hessel","et","al.",",","2017",",","for","a","discussion","on","its","benefits",":","https",":","\/\/arxiv.org\/abs\/1710.02298"]
docstr: Computes loss for DQN training.
  
  Args:
    experience: A batch of experience data in the form of a `Trajectory` or
      `Transition`. The structure of `experience` must match that of
      `self.collect_policy.step_spec`.  If a `Trajectory`, all tensors in
      `experience` must be shaped `[B, T, ...]` where `T` must be equal to
      `self.train_sequence_length` if that property is not `None`.
    td_errors_loss_fn: A function(td_targets, predictions) to compute the
      element wise loss.
    gamma: Discount for future rewards.
    reward_scale_factor: Multiplicative factor to scale rewards.
    weights: Optional scalar or elementwise (per-batch-entry) importance
      weights.  The output td_loss will be scaled by these weights, and the
      final scalar loss is the mean of these values.
    training: Whether this loss is being used for training.
  
  Returns:
    loss: An instance of `DqnLossInfo`.
  Raises:
    ValueError:
      if the number of actions is greater than 1.
code: def _loss(
      self,
      experience,
      td_errors_loss_fn=None,
      gamma=1.0,
      reward_scale_factor=1.0,
      weights=None,
      training=False,
  ):
      """Computes loss for DQN training.
  
      Args:
        experience: A batch of experience data in the form of a `Trajectory` or
          `Transition`. The structure of `experience` must match that of
          `self.collect_policy.step_spec`.  If a `Trajectory`, all tensors in
          `experience` must be shaped `[B, T, ...]` where `T` must be equal to
          `self.train_sequence_length` if that property is not `None`.
        td_errors_loss_fn: A function(td_targets, predictions) to compute the
          element wise loss.
        gamma: Discount for future rewards.
        reward_scale_factor: Multiplicative factor to scale rewards.
        weights: Optional scalar or elementwise (per-batch-entry) importance
          weights.  The output td_loss will be scaled by these weights, and the
          final scalar loss is the mean of these values.
        training: Whether this loss is being used for training.
  
      Returns:
        loss: An instance of `DqnLossInfo`.
      Raises:
        ValueError:
          if the number of actions is greater than 1.
      """
      transition = self._as_transition(experience)
      time_steps, policy_steps, next_time_steps = transition
      actions = policy_steps.action
      # TODO(b/195943557) remove td_errors_loss_fn input to _loss
      self._td_errors_loss_fn = (
          td_errors_loss_fn or self._td_errors_loss_fn
      )
  
      with tf.name_scope("loss"):
          q_values = self._compute_q_values(
              time_steps, actions, training=training
          )
  
          next_q_values = self._compute_next_q_values(
              next_time_steps, policy_steps.info
          )
  
          rewards = reward_scale_factor * next_time_steps.reward
          discounts = gamma * next_time_steps.discount
  
          td_loss, td_error = self._td_loss(
              q_values, next_q_values, rewards, discounts
          )
  
          valid_mask = tf.cast(~time_steps.is_last(), tf.float32)
          td_error = valid_mask * td_error
  
          td_loss = valid_mask * td_loss
  
          if nest_utils.is_batched_nested_tensors(
              time_steps, self.time_step_spec, num_outer_dims=2
          ):
              # Do a sum over the time dimension.
              td_loss = tf.reduce_sum(input_tensor=td_loss, axis=1)
  
          # Aggregate across the elements of the batch and add regularization loss.
          # Note: We use an element wise loss above to ensure each element is always
          #   weighted by 1/N where N is the batch size, even when some of the
          #   weights are zero due to boundary transitions. Weighting by 1/K where K
          #   is the actual number of non-zero weight would artificially increase
          #   their contribution in the loss. Think about what would happen as
          #   the number of boundary samples increases.
  
          agg_loss = common.aggregate_losses(
              per_example_loss=td_loss,
              sample_weight=weights,
              regularization_loss=self._q_network.losses,
          )
          total_loss = agg_loss.total_loss
  
          losses_dict = {
              "td_loss": agg_loss.weighted,
              "reg_loss": agg_loss.regularization,
              "total_loss": total_loss,
          }
  
          common.summarize_scalar_dict(
              losses_dict,
              step=self.train_step_counter,
              name_scope="Losses/",
          )
  
          if self._summarize_grads_and_vars:
              with tf.name_scope("Variables/"):
                  for var in self._q_network.trainable_weights:
                      tf.compat.v2.summary.histogram(
                          name=var.name.replace(":", "_"),
                          data=var,
                          step=self.train_step_counter,
                      )
  
          if self._debug_summaries:
              diff_q_values = q_values - next_q_values
              common.generate_tensor_summaries(
                  "td_error", td_error, self.train_step_counter
              )
              common.generate_tensor_summaries(
                  "td_loss", td_loss, self.train_step_counter
              )
              common.generate_tensor_summaries(
                  "q_values", q_values, self.train_step_counter
              )
              common.generate_tensor_summaries(
                  "next_q_values",
                  next_q_values,
                  self.train_step_counter,
              )
              common.generate_tensor_summaries(
                  "diff_q_values",
                  diff_q_values,
                  self.train_step_counter,
              )
  
          return tf_agent.LossInfo(
              total_loss,
              DqnLossInfo(td_loss=td_loss, td_error=td_error),
          )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  from typing import Optional, Text, cast
  import gin
  import tensorflow
  from tf_agents.agents import data_converter
  from tf_agents.agents import tf_agent
  from tf_agents.networks import network
  from tf_agents.networks import utils
  from tf_agents.policies import boltzmann_policy
  from tf_agents.policies import epsilon_greedy_policy
  from tf_agents.policies import greedy_policy
  from tf_agents.policies import q_policy
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import eager_utils
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 9 for calling the DqnAgent obj's underlying member method code with example usage: obj._loss(experience, td_errors_loss_fn, gamma, reward_scale_factor, weights, training) and returns: tf_agent