code: def make_reverb_dataset(
      server_address: str,
      batch_size: Optional[int] = None,
      prefetch_size: Optional[int] = None,
      table: Union[
          str, Mapping[str, float]
      ] = adders.DEFAULT_PRIORITY_TABLE,
      num_parallel_calls: Optional[int] = 12,
      max_in_flight_samples_per_worker: Optional[int] = None,
      postprocess: Optional[Transform] = None,
      # Deprecated kwargs.
      environment_spec: Optional[specs.EnvironmentSpec] = None,
      extra_spec: Optional[types.NestedSpec] = None,
      transition_adder: bool = False,
      convert_zero_size_to_none: bool = False,
      using_deprecated_adder: bool = False,
      sequence_length: Optional[int] = None,
  ) -> tf.data.Dataset:
      """Make a TensorFlow dataset backed by a Reverb trajectory replay service.
  
      Arguments:
        server_address: Address of the Reverb server.
        batch_size: Batch size of the returned dataset.
        prefetch_size: The number of elements to prefetch from the original dataset.
          Note that Reverb may do some internal prefetching in addition to this.
        table: The name of the Reverb table to use, or a mapping of (table_name,
          float_weight) for mixing multiple tables in the input (e.g. mixing online
          and offline experiences).
        num_parallel_calls: The parralelism to use. Setting it to `tf.data.AUTOTUNE`
          will allow `tf.data` to automatically find a reasonable value.
        max_in_flight_samples_per_worker: see reverb.TrajectoryDataset for details.
        postprocess: User-specified transformation to be applied to the dataset (as
          `ds.map(postprocess)`).
        environment_spec: DEPRECATED! Do not use.
        extra_spec: DEPRECATED! Do not use.
        transition_adder: DEPRECATED! Do not use.
        convert_zero_size_to_none: DEPRECATED! Do not use.
        using_deprecated_adder: DEPRECATED! Do not use.
        sequence_length: DEPRECATED! Do not use.
  
      Returns:
        A `tf.data.Dataset` iterating over the contents of the Reverb table.
  
      Raises:
        ValueError if `environment_spec` or `extra_spec` are set, or `table` is a
        mapping with no positive weight values.
      """
  
      if environment_spec or extra_spec:
          raise ValueError(
              "The make_reverb_dataset factory function no longer requires specs as"
              " as they should be passed as a signature to the reverb.Table when it"
              " is created. Consider either updating your code or falling back to the"
              " deprecated dataset factory in acme/datasets/deprecated."
          )
  
      # These are no longer used and are only kept in the call signature for
      # backward compatibility.
      del environment_spec
      del extra_spec
      del transition_adder
      del convert_zero_size_to_none
      del using_deprecated_adder
      del sequence_length
  
      # This is the default that used to be set by reverb.TFClient.dataset().
      if (
          max_in_flight_samples_per_worker is None
          and batch_size is None
      ):
          max_in_flight_samples_per_worker = 100
      elif max_in_flight_samples_per_worker is None:
          max_in_flight_samples_per_worker = 2 * batch_size
  
      # Create mapping from tables to non-zero weights.
      if isinstance(table, str):
          tables = collections.OrderedDict([(table, 1.0)])
      else:
          tables = collections.OrderedDict(
              [
                  (name, weight)
                  for name, weight in table.items()
                  if weight > 0.0
              ]
          )
          if len(tables) <= 0:
              raise ValueError(
                  f"No positive weights in input tables {tables}"
              )
  
      # Normalize weights.
      total_weight = sum(tables.values())
      tables = collections.OrderedDict(
          [
              (name, weight / total_weight)
              for name, weight in tables.items()
          ]
      )
  
      def _make_dataset(unused_idx: tf.Tensor) -> tf.data.Dataset:
          datasets = ()
          for table_name, weight in tables.items():
              max_in_flight_samples = max(
                  1, int(max_in_flight_samples_per_worker * weight)
              )
              dataset = reverb.TrajectoryDataset.from_table_signature(
                  server_address=server_address,
                  table=table_name,
                  max_in_flight_samples_per_worker=max_in_flight_samples,
              )
              datasets += (dataset,)
          if len(datasets) > 1:
              dataset = tf.data.Dataset.sample_from_datasets(
                  datasets, weights=tables.values()
              )
          else:
              dataset = datasets[0]
  
          # Post-process each element if a post-processing function is passed, e.g.
          # observation-stacking or data augmenting transformations.
          if postprocess:
              dataset = dataset.map(postprocess)
  
          if batch_size:
              dataset = dataset.batch(batch_size, drop_remainder=True)
  
          return dataset
  
      if num_parallel_calls is not None:
          # Create a datasets and interleaves it to create `num_parallel_calls`
          # `TrajectoryDataset`s.
          num_datasets_to_interleave = (
              os.cpu_count()
              if num_parallel_calls == tf.data.AUTOTUNE
              else num_parallel_calls
          )
          dataset = tf.data.Dataset.range(
              num_datasets_to_interleave
          ).interleave(
              map_func=_make_dataset,
              cycle_length=num_parallel_calls,
              num_parallel_calls=num_parallel_calls,
              deterministic=False,
          )
      else:
          dataset = _make_dataset(tf.constant(0))
  
      if prefetch_size:
          dataset = dataset.prefetch(prefetch_size)
  
      return dataset
  
file_path: acme/acme/datasets/reverb.py
imports: import collections
  import os
  from typing import Callable, Mapping, Optional, Union
  from acme import specs
  from acme import types
  from acme.adders import reverb
  import reverb
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: make_reverb_dataset(server_address, batch_size, prefetch_size, table, num_parallel_calls, max_in_flight_samples_per_worker, postprocess, environment_spec, extra_spec, transition_adder, convert_zero_size_to_none, using_deprecated_adder, sequence_length) and returns: dataset, dataset