code: def _generate_samples_with_demonstrations(
      demonstration_iterator: Iterator[types.Transition],
      replay_iterator: Iterator[reverb.ReplaySample],
      policy_to_expert_data_ratio: int,
      batch_size,
  ) -> Iterator[reverb.ReplaySample]:
      """Generator which creates the sample having demonstrations in them.
  
      It takes the demonstrations and replay iterators and generates batches with
      same size as the replay iterator, such that each batches have the ratio of
      policy and expert data specified in policy_to_expert_data_ratio on average.
      There is no constraints on how the demonstrations and replay samples should be
      batched.
  
      Args:
        demonstration_iterator: Iterator of demonstrations.
        replay_iterator: Replay buffer sample iterator.
        policy_to_expert_data_ratio: Amount of policy transitions for 1 expert
          transitions.
        batch_size: Output batch size, which should match the replay batch size.
  
      Yields:
        Samples having a mix of demonstrations and policy data. The info will match
        the current replay sample info and the batch size will be the same as the
        replay_iterator data batch size.
      """
      count = 0
      if batch_size % (policy_to_expert_data_ratio + 1) != 0:
          raise ValueError(
              "policy_to_expert_data_ratio + 1 must divide the batch size but "
              f"{batch_size} % {policy_to_expert_data_ratio+1} !=0"
          )
      demo_insertion_size = batch_size // (
          policy_to_expert_data_ratio + 1
      )
      policy_insertion_size = batch_size - demo_insertion_size
  
      demonstration_iterator = _rebatch(
          demonstration_iterator, demo_insertion_size
      )
      for sample, demos in zip(replay_iterator, demonstration_iterator):
          output_transitions = tree.map_structure(
              functools.partial(
                  _mix_arrays, index=policy_insertion_size, seed=count
              ),
              sample.data,
              demos,
          )
          count += 1
          yield reverb.ReplaySample(
              info=sample.info, data=output_transitions
          )
  
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 a global function with example usage: _generate_samples_with_demonstrations(demonstration_iterator, replay_iterator, policy_to_expert_data_ratio, batch_size) without return types