code: def load_tfrecord_dataset(
      dataset_files,
      buffer_size=1000,
      as_experience=False,
      as_trajectories=False,
      add_batch_dim=True,
      decoder=None,
      num_parallel_reads=None,
      compress_image=False,
      spec=None,
  ):
      """Loads a TFRecord dataset from file, sequencing samples as Trajectories.
  
      Args:
        dataset_files: List of paths to one or more datasets
        buffer_size: (int) number of bytes in the read buffer. 0 means no buffering.
        as_experience: (bool) Returns dataset as a pair of Trajectories. Samples
          will be shaped as if they had been pulled from a replay buffer with
          `num_steps=2`. These samples can be fed directly to agent's `train`
          method.
        as_trajectories: (bool) Remaps the data into trajectory objects. This should
          be enabled when the resulting types must be trajectories as expected by
          agents.
        add_batch_dim: (bool) If True the data will have a batch dim of 1 to conform
          with the expected tensor batch convention. Set to false if you want to
          batch the data on your own.
        decoder: Optional, a custom decoder to use rather than using the default
          spec path.
        num_parallel_reads: Optional, number of parallel reads in the TFRecord
          dataset. If not specified, len(dataset_files) will be used.
        compress_image: Whether to decompress image. It is assumed that any uint8
          tensor of rank 3 with shape (w,h,c) is an image. If the tensor was
          compressed in the encoder, it needs to be decompressed.
        spec: Optional, the dataspec of the TFRecord dataset to be parsed. If not
          provided, parses the dataspec of the TFRecord directly.
  
      Returns:
        A dataset of type tf.data.Dataset. Samples follow the dataset's spec nested
        structure. Samples are generated with a leading batch dim of 1
        (or 2 if as_experience is enabled).
      Raises:
        IOError: One or more of the dataset files does not exist.
      """
  
      if not decoder:
          if spec is None:
              specs = []
              for dataset_file in dataset_files:
                  spec_path = dataset_file + _SPEC_FILE_EXTENSION
                  dataset_spec = parse_encoded_spec_from_file(spec_path)
                  specs.append(dataset_spec)
                  if not all([dataset_spec == spec for spec in specs]):
                      raise IOError(
                          "One or more of the encoding specs do not match."
                      )
              spec = specs[0]
  
          decoder = example_encoding.get_example_decoder(
              spec, compress_image=compress_image
          )
  
      logging.info("Loading TFRecord dataset...")
      if not num_parallel_reads:
          num_parallel_reads = len(dataset_files)
      dataset = tf.data.TFRecordDataset(
          dataset_files,
          buffer_size=buffer_size,
          num_parallel_reads=num_parallel_reads,
      )
  
      def decode_fn(proto):
          """Decodes a proto object."""
          return decoder(proto)
  
      def decode_and_batch_fn(proto):
          """Decodes a proto object, and batch output tensors."""
          sample = decoder(proto)
          return nest_utils.batch_nested_tensors(sample)
  
      if as_experience:
          dataset = dataset.map(
              decode_fn,
              num_parallel_calls=tf.data.experimental.AUTOTUNE,
          ).batch(2, drop_remainder=True)
      elif add_batch_dim:
          dataset = dataset.map(
              decode_and_batch_fn,
              num_parallel_calls=tf.data.experimental.AUTOTUNE,
          )
      else:
          dataset = dataset.map(
              decode_fn,
              num_parallel_calls=tf.data.experimental.AUTOTUNE,
          )
  
      if as_trajectories:
          as_trajectories_fn = lambda sample: trajectory.Trajectory(
              *sample
          )
          dataset = dataset.map(
              as_trajectories_fn,
              num_parallel_calls=tf.data.experimental.AUTOTUNE,
          )
      return dataset
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  from absl import logging
  import tensorflow
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import trajectory
  from tf_agents.utils import eager_utils
  from tf_agents.utils import example_encoding
  from tf_agents.utils import nest_utils
  from tensorflow.core.protobuf import struct_pb2
label_desc: Use image node_id 3 for calling a global function with example usage: load_tfrecord_dataset(dataset_files, buffer_size, as_experience, as_trajectories, add_batch_dim, decoder, num_parallel_reads, compress_image, spec) and returns: dataset, decoder, nest_utils