code: def _make_demonstration_dataset_factory(
      dataset_name: str,
      environment_spec: specs.EnvironmentSpec,
      num_demonstrations: int,
      random_key: jax_types.PRNGKey,
  ) -> Callable[[jax_types.PRNGKey], Iterator[types.Transition]]:
      """Returns the demonstration dataset factory for the given dataset."""
  
      def demonstration_dataset_factory(
          batch_size: int,
      ) -> Iterator[types.Transition]:
          """Returns an iterator of demonstration samples."""
          transitions_iterator = tfds.get_tfds_dataset(
              dataset_name,
              num_episodes=num_demonstrations,
              env_spec=environment_spec,
          )
          return tfds.JaxInMemoryRandomSampleIterator(
              transitions_iterator,
              key=random_key,
              batch_size=batch_size,
          )
  
      return demonstration_dataset_factory
  
file_path: acme/examples/baselines/imitation/run_iqlearn.py
imports: from typing import Callable, Iterator
  from absl import flags
  from acme import specs
  from acme import types
  from acme.agents.jax import actor_core
  from acme.agents.jax import iq_learn
  from acme.datasets import tfds
  import helpers
  from absl import app
  from acme.jax import experiments
  from acme.jax import types
  from acme.utils import lp_utils
  import dm_env
  import jax
  import launchpad
label_desc: Use image node_id 2 for calling a global function with example usage: _make_demonstration_dataset_factory(dataset_name, environment_spec, num_demonstrations, random_key) and returns: demonstration_dataset_factory, tfds