class_docstr: Distributional DQN learner.
class_docstr_tok: ["Distributional","DQN","learner","."]
docstr: Do a step of SGD and update the priorities.
docstr_tok: ["Do","a","step","of","SGD","and","update","the","priorities","."]
code: def _step(self) -> Dict[str, tf.Tensor]:
      """Do a step of SGD and update the priorities."""
  
      # Pull out the data needed for updates/priorities.
      inputs = next(self._iterator)
      transitions: types.Transition = inputs.data
      keys, probs, *_ = inputs.info
  
      with tf.GradientTape() as tape:
          loss, fetches = self._loss_and_fetches(
              transitions.observation,
              transitions.action,
              transitions.reward,
              transitions.discount,
              transitions.next_observation,
          )
  
          # Get the importance weights.
          importance_weights = 1.0 / probs  # [B]
          importance_weights **= self._importance_sampling_exponent
          importance_weights /= tf.reduce_max(importance_weights)
  
          # Reweight.
          loss *= tf.cast(importance_weights, loss.dtype)  # [B]
          loss = tf.reduce_mean(loss, axis=[0])  # []
  
      # Do a step of SGD.
      gradients = tape.gradient(loss, self._network.trainable_variables)
      self._optimizer.apply(
          gradients, self._network.trainable_variables
      )
  
      # Update the priorities in the replay buffer.
      if self._replay_client:
          priorities = tf.clip_by_value(tf.abs(loss), -100, 100)
          priorities = tf.cast(priorities, tf.float64)
          self._replay_client.update_priorities(
              table=adders.DEFAULT_PRIORITY_TABLE,
              keys=keys,
              priorities=priorities,
          )
  
      # Periodically update the target network.
      if tf.math.mod(self._num_steps, self._target_update_period) == 0:
          for src, dest in zip(
              self._network.variables, self._target_network.variables
          ):
              dest.assign(src)
      self._num_steps.assign_add(1)
  
      # Report gradient norms.
      fetches.update(
          loss=loss, gradient_norm=tf.linalg.global_norm(gradients)
      )
      return fetches
  
imports: from typing import Dict, List, Optional, Tuple
  from acme import core
  from acme import types
  from acme.adders import reverb
  from acme.tf import losses
  from acme.tf import networks
  from acme.tf import savers
  from acme.tf import utils
  from acme.utils import counting
  from acme.utils import loggers
  import numpy
  import reverb
  import sonnet
  import tensorflow
label_desc: Use image node_id 2 for calling the IQNLearner obj's underlying member method code with example usage: obj._step() and returns: fetches