class_docstr: R2D2 learner.
  
  This is the learning component of the R2D2 agent. It takes a dataset as input
  and implements update functionality to learn from this dataset.
class_docstr_tok: ["R2D2","learner",".","This","is","the","learning","component","of","the","R2D2","agent",".","It","takes","a","dataset","as","input","and","implements","update","functionality","to","learn","from","this","dataset","."]
code: def _step(self) -> Dict[str, tf.Tensor]:
      # Draw a batch of data from replay.
      sample: reverb.ReplaySample = next(self._iterator)
  
      data = tf2_utils.batch_to_sequence(sample.data)
      observations, actions, rewards, discounts, extra = (
          data.observation,
          data.action,
          data.reward,
          data.discount,
          data.extras,
      )
      unused_sequence_length, batch_size = actions.shape
  
      # Get initial state for the LSTM, either from replay or simply use zeros.
      if self._store_lstm_state:
          core_state = tree.map_structure(
              lambda x: x[0], extra["core_state"]
          )
      else:
          core_state = self._network.initial_state(batch_size)
      target_core_state = tree.map_structure(tf.identity, core_state)
  
      # Before training, optionally unroll the LSTM for a fixed warmup period.
      burn_in_obs = tree.map_structure(
          lambda x: x[: self._burn_in_length], observations
      )
      _, core_state = self._burn_in(burn_in_obs, core_state)
      _, target_core_state = self._burn_in(
          burn_in_obs, target_core_state
      )
  
      # Don't train on the warmup period.
      (
          observations,
          actions,
          rewards,
          discounts,
          extra,
      ) = tree.map_structure(
          lambda x: x[self._burn_in_length :],
          (observations, actions, rewards, discounts, extra),
      )
  
      with tf.GradientTape() as tape:
          # Unroll the online and target Q-networks on the sequences.
          q_values, _ = self._network.unroll(
              observations, core_state, self._sequence_length
          )
          target_q_values, _ = self._target_network.unroll(
              observations, target_core_state, self._sequence_length
          )
  
          # Compute the target policy distribution (greedy).
          greedy_actions = tf.argmax(
              q_values, output_type=tf.int32, axis=-1
          )
          target_policy_probs = tf.one_hot(
              greedy_actions,
              depth=self._num_actions,
              dtype=q_values.dtype,
          )
  
          # Compute the transformed n-step loss.
          rewards = tree.map_structure(lambda x: x[:-1], rewards)
          discounts = tree.map_structure(lambda x: x[:-1], discounts)
          loss, extra = losses.transformed_n_step_loss(
              qs=q_values,
              targnet_qs=target_q_values,
              actions=actions,
              rewards=rewards,
              pcontinues=discounts * self._discount,
              target_policy_probs=target_policy_probs,
              bootstrap_n=self._n_step,
          )
  
          # Calculate importance weights and use them to scale the loss.
          sample_info = sample.info
          keys, probs = sample_info.key, sample_info.probability
          importance_weights = 1.0 / (
              self._max_replay_size * probs
          )  # [T, B]
          importance_weights **= self._importance_sampling_exponent
          importance_weights /= tf.reduce_max(importance_weights)
          loss *= tf.cast(importance_weights, tf.float32)  # [T, B]
          loss = tf.reduce_mean(loss)  # []
  
      # Apply gradients via optimizer.
      gradients = tape.gradient(loss, self._network.trainable_variables)
      # Clip and apply gradients.
      if self._clip_grad_norm is not None:
          gradients, _ = tf.clip_by_global_norm(
              gradients, self._clip_grad_norm
          )
  
      self._optimizer.apply(
          gradients, self._network.trainable_variables
      )
  
      # 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)
  
      if self._reverb_client:
          # Compute updated priorities.
          priorities = compute_priority(
              extra.errors, self._max_priority_weight
          )
          # Compute priorities and add an op to update them on the reverb side.
          self._reverb_client.update_priorities(
              table=adders.DEFAULT_PRIORITY_TABLE,
              keys=keys,
              priorities=tf.cast(priorities, tf.float64),
          )
  
      return {"loss": loss}
  
imports: import functools
  import time
  from typing import Dict, Iterator, List, Mapping, Union, Optional
  import acme
  from acme import specs
  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
  import tree
label_desc: Use image node_id 2 for calling the R2D2Learner obj's underlying member method code with example usage: obj._step() and returns: dict