class_docstr: D4PG learner.
  
  This is the learning component of a D4PG agent. IE it takes a dataset as input
  and implements update functionality to learn from this dataset.
class_docstr_tok: ["D4PG","learner",".","This","is","the","learning","component","of","a","D4PG","agent",".","IE","it","takes","a","dataset","as","input","and","implements","update","functionality","to","learn","from","this","dataset","."]
code: def _step(self, sample) -> Dict[str, tf.Tensor]:
      transitions: types.Transition = (
          sample.data
      )  # Assuming ReverbSample.
  
      # Cast the additional discount to match the environment discount dtype.
      discount = tf.cast(
          self._discount, dtype=transitions.discount.dtype
      )
  
      with tf.GradientTape(persistent=True) as tape:
          # Maybe transform the observation before feeding into policy and critic.
          # Transforming the observations this way at the start of the learning
          # step effectively means that the policy and critic share observation
          # network weights.
          o_tm1 = self._observation_network(transitions.observation)
          o_t = self._target_observation_network(
              transitions.next_observation
          )
          # This stop_gradient prevents gradients to propagate into the target
          # observation network. In addition, since the online policy network is
          # evaluated at o_t, this also means the policy loss does not influence
          # the observation network training.
          o_t = tree.map_structure(tf.stop_gradient, o_t)
  
          # Critic learning.
          q_tm1 = self._critic_network(o_tm1, transitions.action)
          q_t = self._target_critic_network(
              o_t, self._target_policy_network(o_t)
          )
  
          # Critic loss.
          critic_loss = losses.categorical(
              q_tm1,
              transitions.reward,
              discount * transitions.discount,
              q_t,
          )
          critic_loss = tf.reduce_mean(critic_loss, axis=[0])
  
          # Actor learning.
          dpg_a_t = self._policy_network(o_t)
          dpg_z_t = self._critic_network(o_t, dpg_a_t)
          dpg_q_t = dpg_z_t.mean()
  
          # Actor loss. If clipping is true use dqda clipping and clip the norm.
          dqda_clipping = 1.0 if self._clipping else None
          policy_loss = losses.dpg(
              dpg_q_t,
              dpg_a_t,
              tape=tape,
              dqda_clipping=dqda_clipping,
              clip_norm=self._clipping,
          )
          policy_loss = tf.reduce_mean(policy_loss, axis=[0])
  
      # Get trainable variables.
      policy_variables = self._policy_network.trainable_variables
      critic_variables = (
          # In this agent, the critic loss trains the observation network.
          self._observation_network.trainable_variables
          + self._critic_network.trainable_variables
      )
  
      # Compute gradients.
      replica_context = tf.distribute.get_replica_context()
      policy_gradients = _average_gradients_across_replicas(
          replica_context, tape.gradient(policy_loss, policy_variables)
      )
      critic_gradients = _average_gradients_across_replicas(
          replica_context, tape.gradient(critic_loss, critic_variables)
      )
  
      # Delete the tape manually because of the persistent=True flag.
      del tape
  
      # Maybe clip gradients.
      if self._clipping:
          policy_gradients = tf.clip_by_global_norm(
              policy_gradients, 40.0
          )[0]
          critic_gradients = tf.clip_by_global_norm(
              critic_gradients, 40.0
          )[0]
  
      # Apply gradients.
      self._policy_optimizer.apply(policy_gradients, policy_variables)
      self._critic_optimizer.apply(critic_gradients, critic_variables)
  
      # Losses to track.
      return {
          "critic_loss": critic_loss,
          "policy_loss": policy_loss,
      }
  
imports: import time
  from typing import Dict, Iterator, List, Optional, Union, Sequence
  import acme
  from acme import types
  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 D4PGLearner obj's underlying member method code with example usage: obj._step(sample) and returns: dict