code: def make_demonstrations(
      env: dm_env.Environment, batch_size: int
  ) -> tf.data.Dataset:
      """Prepare the dataset of demonstrations."""
      batch_dataset = bsuite_demonstrations.make_dataset(
          env, stochastic=False
      )
      # Combine with demonstration dataset.
      transition = functools.partial(
          _n_step_transition_from_episode,
          n_step=1,
          additional_discount=1.0,
      )
  
      dataset = batch_dataset.map(transition)
  
      # Batch and prefetch.
      dataset = dataset.batch(batch_size, drop_remainder=True)
      dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
  
      return dataset
  
imports: import functools
  import operator
  from typing import Callable
  from acme import core
  from acme import environment_loop
  from acme import specs
  from acme import types
  from acme.agents.jax import actor_core
  from acme.agents.jax import actors
  from acme.agents.jax import bc
  from acme.agents.tf.dqfd import bsuite_demonstrations
  from acme.jax import networks
  from acme.jax import types
  from acme.jax import utils
  from acme.jax import variable_utils
  from acme.jax.deprecated import offline_distributed_layout
  from acme.utils import counting
  from acme.utils import loggers
  from acme.wrappers import single_precision
  import bsuite
  import dm_env
  import haiku
  import jax
  import jax.numpy
  from jax.scipy import special
  import rlax
  import tensorflow
  import tree
label_desc: Use image node_id 4 for calling a global function with example usage: make_demonstrations(env, batch_size) and returns: dataset