class_docstr: AIL learner.
class_docstr_tok: ["AIL","learner","."]
code: def step(self):
      sample = next(self._iterator)
      rb_transitions = sample.discriminator_sample
      demo_transitions = sample.demonstration_sample
  
      if demo_transitions.reward.shape != rb_transitions.reward.shape:
          raise ValueError(
              "Different shapes for demo transitions and rb_transitions: "
              f"{demo_transitions.reward.shape} != {rb_transitions.reward.shape}"
          )
  
      # Update the parameters of the policy before doing a gradient step.
      state = self._state._replace(
          policy_params=self._get_policy_params()
      )
      self._state, metrics = self._update_step(
          state, (demo_transitions, rb_transitions)
      )
  
      # The order is important for AIRL.
      # In AIRL, the discriminator update depends on the logpi of the direct rl
      # policy.
      # When updating the discriminator, we want the logpi for which the
      # transitions were made with and not an updated one.
      # Get data from replay (dropping extras if any). Note there is no
      # extra data here because we do not insert any into Reverb.
      self._direct_rl_learner.step()
  
      # 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 functools
  import itertools
  import time
  from typing import Any, Callable, Iterator, List, NamedTuple, Optional, Tuple
  import acme
  from acme import types
  from acme.agents.jax.ail import losses
  from acme.agents.jax.ail import networks
  from acme.jax import networks
  from acme.jax import utils
  from acme.utils import counting
  from acme.utils import loggers
  from acme.utils import reverb_utils
  import jax
  import optax
  import reverb
label_desc: Use image node_id 3 for calling the AILLearner obj's underlying member method code with example usage: obj.step() without return types