class_docstr: Learner for an importanced-weighted advantage actor-critic.
class_docstr_tok: ["Learner","for","an","importanced-weighted","advantage","actor-critic","."]
docstr: Does a step of SGD and logs the results.
docstr_tok: ["Does","a","step","of","SGD","and","logs","the","results","."]
code: def step(self):
      """Does a step of SGD and logs the results."""
      samples = next(self._iterator)
  
      # Do a batch of SGD.
      start = time.time()
      self._state, results = self._sgd_step(self._state, samples)
  
      # Take results from first replica.
      # NOTE: This measure will be a noisy estimate for the purposes of the logs
      # as it does not pmean over all devices.
      results = utils.get_from_first_device(results)
  
      # Update our counts and record them.
      counts = self._counter.increment(
          steps=1, time_elapsed=time.time() - start
      )
  
      # Maybe write logs.
      self._logger.write({**results, **counts})
  
imports: import time
  from typing import Dict, Iterator, List, NamedTuple, Optional, Sequence, Tuple
  from absl import logging
  import acme
  from acme.agents.jax.impala import networks
  from acme.jax import losses
  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 numpy
  import optax
  import reverb
label_desc: Use image node_id 2 for calling the IMPALALearner obj's underlying member method code with example usage: obj.step() without return types