class_docstr: AIL Builder.
class_docstr_tok: ["AIL","Builder","."]
code: def make_dataset_iterator(
      self, replay_client: reverb.Client
  ) -> Iterator[learning.AILSample]:
      batch_size_per_learner_step = (
          ail_config.get_per_learner_step_batch_size(self._config)
      )
  
      iterator_demonstration = self._make_demonstrations(
          batch_size_per_learner_step
      )
  
      direct_iterator = self._rl_agent.make_dataset_iterator(
          replay_client
      )
  
      if self._config.share_iterator:
          # In order to reuse the iterator return values and not lose a 2x factor on
          # sample efficiency, we need to use itertools.tee().
          discriminator_iterator, direct_iterator = itertools.tee(
              direct_iterator
          )
      else:
          discriminator_iterator = datasets.make_reverb_dataset(
              table=self._config.replay_table_name,
              server_address=replay_client.server_address,
              batch_size=ail_config.get_per_learner_step_batch_size(
                  self._config
              ),
              prefetch_size=self._config.prefetch_size,
          ).as_numpy_iterator()
  
      if self._config.policy_to_expert_data_ratio is not None:
          (
              iterator_demonstration,
              iterator_demonstration2,
          ) = itertools.tee(iterator_demonstration)
          direct_iterator = _generate_samples_with_demonstrations(
              iterator_demonstration2,
              direct_iterator,
              self._config.policy_to_expert_data_ratio,
              self._config.direct_rl_batch_size,
          )
  
      is_sequence_based = self._config.is_sequence_based
  
      # Don't flatten the discriminator batch if the iterator is not shared.
      process_discriminator_sample = functools.partial(
          reverb_utils.replay_sample_to_sars_transition,
          is_sequence=is_sequence_based and self._config.share_iterator,
          flatten_batch=is_sequence_based
          and self._config.share_iterator,
          strip_last_transition=is_sequence_based
          and self._config.share_iterator,
      )
  
      discriminator_iterator = (
          # Remove the extras to have the same nested structure as demonstrations.
          process_discriminator_sample(sample)._replace(extras=())
          for sample in discriminator_iterator
      )
  
      return utils.device_put(
          (
              learning.AILSample(*sample)
              for sample in zip(
                  discriminator_iterator,
                  direct_iterator,
                  iterator_demonstration,
              )
          ),
          jax.devices()[0],
      )
  
imports: import functools
  import itertools
  from typing import Callable, Generic, Iterator, List, Optional, Tuple
  from acme import adders
  from acme import core
  from acme import specs
  from acme import types
  from acme.adders import reverb
  from acme.agents.jax import builders
  from acme.agents.jax.ail import config
  from acme.agents.jax.ail import learning
  from acme.agents.jax.ail import losses
  from acme.agents.jax.ail import networks
  from acme.datasets import reverb
  from acme.jax import types
  from acme.jax import utils
  from acme.jax.imitation_learning_types import DirectPolicyNetwork
  from acme.utils import counting
  from acme.utils import loggers
  from acme.utils import reverb_utils
  import jax
  import numpy
  import optax
  import reverb
  from reverb import rate_limiters
  import tree
label_desc: Use image node_id 4 for calling the AILBuilder obj's underlying member method code with example usage: obj.make_dataset_iterator(replay_client) and returns: utils