code: def _generate_sqil_samples(
      demonstration_iterator: Iterator[types.Transition],
      replay_iterator: Iterator[reverb.ReplaySample],
  ) -> Iterator[reverb.ReplaySample]:
      """Generator which creates the sample iterator for SQIL.
  
      Args:
        demonstration_iterator: Iterator of demonstrations.
        replay_iterator: Replay buffer sample iterator.
  
      Yields:
        Samples having a mix of demonstrations with reward 1 and replay samples with
        reward 0.
      """
      for demonstrations, replay_sample in zip(
          demonstration_iterator, replay_iterator
      ):
          demonstrations = demonstrations._replace(
              reward=np.ones_like(demonstrations.reward)
          )
  
          replay_transitions = replay_sample.data
          replay_transitions = replay_transitions._replace(
              reward=np.zeros_like(replay_transitions.reward)
          )
  
          double_batch = tree.map_structure(
              lambda x, y: np.concatenate([x, y]),
              demonstrations,
              replay_transitions,
          )
  
          # Split the double batch in an interleaving fashion.
          # e.g [1, 2, 3, 4 ,5 ,6] -> [1, 3, 5] and [2, 4, 6]
          yield reverb.ReplaySample(
              info=replay_sample.info,
              data=tree.map_structure(lambda x: x[0::2], double_batch),
          )
          yield reverb.ReplaySample(
              info=replay_sample.info,
              data=tree.map_structure(lambda x: x[1::2], double_batch),
          )
  
imports: from typing import Callable, Generic, Iterator, List, Optional
  from acme import adders
  from acme import core
  from acme import specs
  from acme import types
  from acme.agents.jax import builders
  from acme.jax import networks
  from acme.jax import utils
  from acme.jax.imitation_learning_types import DirectPolicyNetwork, DirectRLNetworks
  from acme.utils import counting
  from acme.utils import loggers
  import jax
  import numpy
  import reverb
  import tree
label_desc: Use image node_id 1 for calling a global function with example usage: _generate_sqil_samples(demonstration_iterator, replay_iterator) without return types