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.
code: def step(self):
      # Sample from replay and pack the data in a Transition.
      sample = next(self._iterator)
      transitions = types.Transition(*sample.data)
  
      self._state, metrics = self._sgd_step(self._state, transitions)
  
      # Take the metrics from the first device, since they've been pmeaned over
      # all devices and are therefore identical.
      metrics = utils.get_from_first_device(metrics)
  
      # Compute elapsed time.
      timestamp = time.time()
      elapsed_time = (
          timestamp - self._timestamp if self._timestamp else 0
      )
      self._timestamp = timestamp
  
      # Increment counts and record the current time
      counts = self._counter.increment(steps=1, walltime=elapsed_time)
  
      # Attempts to write the logs.
      self._logger.write({**metrics, **counts})
  
imports: import time
  from typing import Dict, Iterator, List, NamedTuple, Optional, Tuple
  import acme
  from acme import types
  from acme.jax import networks
  from acme.jax import utils
  from acme.utils import counting
  from acme.utils import loggers
  import jax
  import jax.numpy
  import optax
  import reverb
  import rlax
label_desc: Use image node_id 2 for calling the D4PGLearner obj's underlying member method code with example usage: obj.step() without return types