Function read_d4rl_dataset

tensorflow_datasets/d4rl/dataset_utils.py:260–272

Reads a D4RL dataset and returns the dataset as a dictionary.

(file_path: str)

def read_d4rl_dataset(file_path: str):
  """Reads a D4RL dataset and returns the dataset as a dictionary."""
  with tf.io.gfile.GFile(file_path, 'rb') as f:
    with h5py.File(f, 'r') as dataset_file:
      dataset_dict = {}
      for k in _get_dataset_keys(dataset_file):
        try:
          # first try loading as an array
          dataset_dict[k] = dataset_file[k][:]
        except ValueError:  # try loading as a scalar
          dataset_dict[k] = dataset_file[k][()]

    return dataset_dict