class_docstr: AIL learner.
class_docstr_tok: ["AIL","learner","."]
docstr: AIL Learner.
  
  Args:
    counter: Counter.
    direct_rl_learner_factory: Function that creates the direct RL learner
      when passed a replay sample iterator.
    loss_fn: Discriminator loss.
    iterator: Iterator that returns AILSamples.
    discriminator_optimizer: Discriminator optax optimizer.
    ail_network: AIL networks.
    discriminator_key: RNG key.
    is_sequence_based: If True, a direct rl algorithm is using SequenceAdder
      data format. Otherwise the learner assumes that the direct rl algorithm
      is using NStepTransitionAdder.
    num_sgd_steps_per_step: Number of discriminator gradient updates per step.
    policy_variable_name: The name of the policy variable to retrieve
      direct_rl policy parameters.
    logger: Logger.
docstr_tok: ["AIL","Learner",".","Args",":","counter",":","Counter",".","direct_rl_learner_factory",":","Function","that","creates","the","direct","RL","learner","when","passed","a","replay","sample","iterator",".","loss_fn",":","Discriminator","loss",".","iterator",":","Iterator","that","returns","AILSamples",".","discriminator_optimizer",":","Discriminator","optax","optimizer",".","ail_network",":","AIL","networks",".","discriminator_key",":","RNG","key",".","is_sequence_based",":","If","True",",","a","direct","rl","algorithm","is","using","SequenceAdder","data","format",".","Otherwise","the","learner","assumes","that","the","direct","rl","algorithm","is","using","NStepTransitionAdder",".","num_sgd_steps_per_step",":","Number","of","discriminator","gradient","updates","per","step",".","policy_variable_name",":","The","name","of","the","policy","variable","to","retrieve","direct_rl","policy","parameters",".","logger",":","Logger","."]
code: def __init__(
      self,
      counter: counting.Counter,
      direct_rl_learner_factory: Callable[
          [Iterator[reverb.ReplaySample]], acme.Learner
      ],
      loss_fn: losses.Loss,
      iterator: Iterator[AILSample],
      discriminator_optimizer: optax.GradientTransformation,
      ail_network: ail_networks.AILNetworks,
      discriminator_key: networks_lib.PRNGKey,
      is_sequence_based: bool,
      num_sgd_steps_per_step: int = 1,
      policy_variable_name: Optional[str] = None,
      logger: Optional[loggers.Logger] = None,
  ):
      """AIL Learner.
  
      Args:
        counter: Counter.
        direct_rl_learner_factory: Function that creates the direct RL learner
          when passed a replay sample iterator.
        loss_fn: Discriminator loss.
        iterator: Iterator that returns AILSamples.
        discriminator_optimizer: Discriminator optax optimizer.
        ail_network: AIL networks.
        discriminator_key: RNG key.
        is_sequence_based: If True, a direct rl algorithm is using SequenceAdder
          data format. Otherwise the learner assumes that the direct rl algorithm
          is using NStepTransitionAdder.
        num_sgd_steps_per_step: Number of discriminator gradient updates per step.
        policy_variable_name: The name of the policy variable to retrieve
          direct_rl policy parameters.
        logger: Logger.
      """
      self._is_sequence_based = is_sequence_based
  
      state_key, networks_key = jax.random.split(discriminator_key)
  
      # Generator expression that works the same as an iterator.
      # https://pymbook.readthedocs.io/en/latest/igd.html#generator-expressions
      iterator, direct_rl_iterator = itertools.tee(iterator)
      direct_rl_iterator = (
          self._process_sample(sample.direct_sample)
          for sample in direct_rl_iterator
      )
      self._direct_rl_learner = direct_rl_learner_factory(
          direct_rl_iterator
      )
  
      self._iterator = iterator
  
      if policy_variable_name is not None:
  
          def get_policy_params():
              return self._direct_rl_learner.get_variables(
                  [policy_variable_name]
              )[0]
  
          self._get_policy_params = get_policy_params
  
      else:
          self._get_policy_params = lambda: None
  
      # General learner book-keeping and loggers.
      self._counter = counter or counting.Counter()
      self._logger = logger or loggers.make_default_logger(
          "learner",
          asynchronous=True,
          serialize_fn=utils.fetch_devicearray,
          steps_key=self._counter.get_steps_key(),
      )
  
      # Use the JIT compiler.
      self._update_step = functools.partial(
          ail_update_step,
          optimizer=discriminator_optimizer,
          ail_network=ail_network,
          loss_fn=loss_fn,
      )
      self._update_step = utils.process_multiple_batches(
          self._update_step, num_sgd_steps_per_step
      )
      self._update_step = jax.jit(self._update_step)
  
      (
          discriminator_params,
          discriminator_state,
      ) = ail_network.discriminator_network.init(networks_key)
      self._state = DiscriminatorTrainingState(
          optimizer_state=discriminator_optimizer.init(
              discriminator_params
          ),
          discriminator_params=discriminator_params,
          discriminator_state=discriminator_state,
          policy_params=self._get_policy_params(),
          key=state_key,
          steps=0,
      )
  
      # Do not record timestamps until after the first learning step is done.
      # This is to avoid including the time it takes for actors to come online and
      # fill the replay buffer.
      self._timestamp = None
  
      self._get_reward = jax.jit(
          functools.partial(
              ail_networks.compute_ail_reward, networks=ail_network
          )
      )
  
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 1 to create a new AILLearner object from inherited base classes: acme with example: obj = AILLearner(counter, direct_rl_learner_factory, loss_fn, iterator, discriminator_optimizer, ail_network, discriminator_key, is_sequence_based, num_sgd_steps_per_step, policy_variable_name, logger)