class_docstr: Learner for an importanced-weighted advantage actor-critic.
docstr: Does an SGD step on a batch of sequences.
docstr_tok: ["Does","an","SGD","step","on","a","batch","of","sequences","."]
code: def _step(self) -> Dict[str, tf.Tensor]:
      """Does an SGD step on a batch of sequences."""
  
      # Retrieve a batch of data from replay.
      inputs: reverb.ReplaySample = next(self._iterator)
      data = tf2_utils.batch_to_sequence(inputs.data)
      observations, actions, rewards, discounts, extra = (
          data.observation,
          data.action,
          data.reward,
          data.discount,
          data.extras,
      )
      core_state = tree.map_structure(
          lambda s: s[0], extra["core_state"]
      )
  
      #
      actions = actions[:-1]  # [T-1]
      rewards = rewards[:-1]  # [T-1]
      discounts = discounts[:-1]  # [T-1]
  
      with tf.GradientTape() as tape:
          # Unroll current policy over observations.
          (logits, values), _ = snt.static_unroll(
              self._network, observations, core_state
          )
  
          # Compute importance sampling weights: current policy / behavior policy.
          behaviour_logits = extra["logits"]
          pi_behaviour = tfd.Categorical(logits=behaviour_logits[:-1])
          pi_target = tfd.Categorical(logits=logits[:-1])
          log_rhos = pi_target.log_prob(
              actions
          ) - pi_behaviour.log_prob(actions)
  
          # Optionally clip rewards.
          rewards = tf.clip_by_value(
              rewards,
              tf.cast(-self._max_abs_reward, rewards.dtype),
              tf.cast(self._max_abs_reward, rewards.dtype),
          )
  
          # Critic loss.
          vtrace_returns = trfl.vtrace_from_importance_weights(
              log_rhos=tf.cast(log_rhos, tf.float32),
              discounts=tf.cast(self._discount * discounts, tf.float32),
              rewards=tf.cast(rewards, tf.float32),
              values=tf.cast(values[:-1], tf.float32),
              bootstrap_value=values[-1],
          )
          critic_loss = tf.square(vtrace_returns.vs - values[:-1])
  
          # Policy-gradient loss.
          policy_gradient_loss = trfl.policy_gradient(
              policies=pi_target,
              actions=actions,
              action_values=vtrace_returns.pg_advantages,
          )
  
          # Entropy regulariser.
          entropy_loss = trfl.policy_entropy_loss(pi_target).loss
  
          # Combine weighted sum of actor & critic losses.
          loss = tf.reduce_mean(
              policy_gradient_loss
              + self._baseline_cost * critic_loss
              + self._entropy_cost * entropy_loss
          )
  
      # Compute gradients and optionally apply clipping.
      gradients = tape.gradient(loss, self._network.trainable_variables)
      gradients, _ = tf.clip_by_global_norm(
          gradients, self._max_gradient_norm
      )
      self._optimizer.apply(
          gradients, self._network.trainable_variables
      )
  
      metrics = {
          "loss": loss,
          "critic_loss": tf.reduce_mean(critic_loss),
          "entropy_loss": tf.reduce_mean(entropy_loss),
          "policy_gradient_loss": tf.reduce_mean(policy_gradient_loss),
      }
  
      return metrics
  
imports: import time
  from typing import Dict, List, Mapping, Optional
  import acme
  from acme import specs
  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 tensorflow_probability
  import tree
  import trfl
label_desc: Use image node_id 3 for calling the IMPALALearner obj's underlying member method code with example usage: obj._step() and returns: metrics