class_docstr: Internal BCQ learner.
  
  This implements the Q-learning component in the discrete BCQ algorithm.
class_docstr_tok: ["Internal","BCQ","learner",".","This","implements","the","Q-learning","component","in","the","discrete","BCQ","algorithm","."]
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[:2]
  
      with tf.GradientTape() as tape:
          # Evaluate our networks.
          q_tm1 = self._q_network(transitions.observation)
          q_t_value = self._target_q_network(
              transitions.next_observation
          )
          q_t_selector = self._network(transitions.next_observation)
  
          # The rewards and discounts have to have the same type as network values.
          r_t = tf.cast(transitions.reward, q_tm1.dtype)
          r_t = tf.clip_by_value(r_t, -1.0, 1.0)
          d_t = tf.cast(transitions.discount, q_tm1.dtype) * tf.cast(
              self._discount, q_tm1.dtype
          )
  
          # Compute the loss.
          _, extra = trfl.double_qlearning(
              q_tm1,
              transitions.action,
              r_t,
              d_t,
              q_t_value,
              q_t_selector,
          )
          loss = losses.huber(
              extra.td_error, self._huber_loss_parameter
          )
  
          # 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.cast(tf.abs(extra.td_error), 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._q_network.variables,
              self._target_q_network.variables,
          ):
              dest.assign(src)
      self._num_steps.assign_add(1)
  
      # Compute the global norm of the gradients for logging.
      global_gradient_norm = tf.linalg.global_norm(gradients)
  
      # Compute statistics of the Q-values for logging.
      max_q = tf.reduce_max(q_t_value)
      min_q = tf.reduce_min(q_t_value)
      mean_q, var_q = tf.nn.moments(q_t_value, [0, 1])
  
      # Report loss & statistics for logging.
      fetches = {
          "gradient_norm": global_gradient_norm,
          "loss": loss,
          "max_q": max_q,
          "mean_q": mean_q,
          "min_q": min_q,
          "var_q": var_q,
      }
  
      return fetches
  
imports: import copy
  from typing import Dict, List, Optional
  from acme import core
  from acme import types
  from acme.adders import reverb
  from acme.agents.tf import bc
  from acme.tf import losses
  from acme.tf import savers
  from acme.tf import utils
  from acme.tf.networks import discrete
  from acme.utils import counting
  from acme.utils import loggers
  import numpy
  import reverb
  import sonnet
  import tensorflow
  import trfl
label_desc: Use image node_id 2 for calling the _InternalBCQLearner obj's underlying member method code with example usage: obj._step() and returns: fetches