Home

Awesome

RLDS

RLDS stands for Reinforcement Learning Datasets and it is an ecosystem of tools to store, retrieve and manipulate episodic data in the context of Sequential Decision Making including Reinforcement Learning (RL), Learning for Demonstrations, Offline RL or Imitation Learning.

This repository includes a library for manipulating RLDS compliant datasets. For other parts of the pipeline please refer to:

Learn more about the RLDS ecosystem in the Google AI Blog and the arXiv paper.

QuickStart & Colabs

See how to use RLDS in this tutorial.

You can find more examples, including performance best practices in the examples page. Besides, the transformations page provides an overview of the RLDS library.

Available datasets

This is a non-exhaustive list of datasets that are compatible with RLDS:

If you want to add your dataset to this list, let us know!

Dataset Format

The dataset is retrieved as a tf.data.Dataset of Episodes where each episode contains a tf.data.Dataset of steps.

drawing

How to create a dataset

Although you can read datasets with the RLDS format even if they were not created with our tools (for example, by adding them to TFDS), we recommend the use of EnvLogger and RLDS Creator as they ensure that the data is stored in a lossless fashion and compatible with RLDS.

Synthetic datasets

Envlogger provides a dm_env Environment class wrapper that records interactions between a real environment and an agent.

env = envlogger.EnvLogger(
      environment,
      data_directory=`/tmp/mydataset`)

Besides, two callbacks can be passed to the EnvLogger constructor to store per-step metadata and per-episode metadata. See the EnvLogger documentation for more details.

Note that per-session metadata can be stored but is currently ignored when loading the dataset.

NOTE: We recommend to use the TFDS Envlogger backend in order to get datasets that can be read directly with TFDS. See an example in this colab.

Note that the Envlogger follows the dm_env convention. So considering:

Data is generated as:

    (o_0, _, _, _, m_0) → (o_1, a_0, r_0, d_0, m_1)  → (o_2, a_1, r_1, d_1, m_2) ⇢ ...

But loaded with RLDS as:

    (o_0,a_0, r_0, d_0, m_0) → (o_1, a_1, r_1, d_1, m_1)  → (o_2, a_2, r_2, d_2, m_2) ⇢ ...

Human datasets

If you want to collect data generated by a human interacting with an environment, check the RLDS Creator.

How to load a dataset

RL datasets can be loaded with TFDS and they are retrieved with the canonical RLDS dataset format.

Load with TFDS

Note: In TFDS you can load the nested dataset as a batched sequence instead of a tf.data.Dataset. See the FAQ for details.

Datasets created with Envlogger and the TFDS backend

These datasets can be loaded directly with:

tfds.builder_from_directory('path').as_dataset(split='all')

or from a list of paths:

tfds.builder_from_directories(paths).as_dataset(split='all')

See more examples in this colab.

Datasets in the TFDS catalog

These datasets can be loaded directly with:

tfds.load('dataset_name').as_dataset()['train']

This is how we load the datasets in the tutorial.

See the full documentation and the catalog in the [TFDS] site.

Datasets in your own repository

Datasets can be implemented with TFDS both inside and outside of the TFDS repository. See examples here.

How to add your dataset to TFDS

This is only necessary when your dataset is not already in TFDS format or if you want to add it to the TFDS catalog. See more details in this page.

Performance best practices

As RLDS exposes RL datasets in a form of Tensorflow's tf.data, many Tensorflow's performance hints apply to RLDS as well. It is important to note, however, that RLDS datasets are very specific and not all general speed-up methods work out of the box. Advice on improving performance might not result in expected outcome.

RLDS provides an optimized library of transformations, but to get a better understanding on how to use RLDS datasets effectively we recommend going through this colab.

FAQ

Processing steps in random order

While by default the order of episodes in RLDS datasets is randomized and there is no need to randomize them again when loading the dataset, some algorithms operate on steps/n-step transitions. There are different ways to interleave steps across multiple episodes - for example:

def ds_loader():
  episode_dataset = tfds.load(...)
  step_dataset = episode_dataset.flat_map(lambda x: x[rlds.STEPS])
  return step_dataset

dataset = Dataset.range(1, N).interleave(ds_loader, cycle_length=..., block_length=...)

Each copy of the dataset shuffles input partitions independently, so consecutive steps returned by the resulting dataset come from unrelated episodes. It is important to note, however, that this way each step will be loaded N times. To avoid duplicates, it is possible to construct each dataset using disjoint splits.

See one example of randomized access in the Atari colab.

Processing random episodes in multiple readers.

Sometimes, users read multiple copies of the dataset in separate processes. For example, to emulate a multiple-actor single learner scenario, where the actors get the offline data from the same dataset. In these situations, it is important that the different processes don't get the same sequence of episodes.

When the number of readers is known, the easiest way is to use the split API from TFDS to ensure that each of the reader takes a different set of episodes from the dataset. Note that if one of the reader dies, its portion of the dataset will not be processed.

Another option is to ensure that the datasets are read in a non-deterministic way. This can be achieved by setting shuffle_files=True and by tuning the ReadConfig options in tfds.load or in builder.as_dataset. You can find more details in the [TFDS documentation about determinism]. In this case, if a reader dies, the full dataset can still be processed. However, with this option, some episodes may appear more than once.

Reducing memory usage

To improve throughput of loading datasets, by default TFDS loads multiple partitions of the dataset in parallel. In the case of datasets with big episodes that can result in high memory usage. If you run into high memory usage problems, it is worth playing around with read_config provided to tfds.load.

Loading the steps as a batch instead of a nested dataset.

If using TFDS you can load the nested dataset as a batched sequence instead of a nested tf.data.Dataset. You can do it by using SkipDecoding:

ds = tfds.load('d4rl_mujoco_halfcheetah/v0-medium', decoders={rlds.STEPS: tfds.decode.SkipDecoding()}, split='train')

To decode the steps as a dataset, you can use tf.data.Dataset.from_tensor_slices.


for e in ds:
 print(tf.data.Dataset.from_tensor_slices(e[rlds.STEPS]))
 break

When using tfds.builder_from_directories or tfds.builder_from_directory, the decoder argument can be passed to as_dataset.

Who uses RLDS

Publications

Below is a sample of publications using RLDS:

Citation

If you use RLDS, please cite the RLDS paper as

@misc{ramos2021rlds,
      title={RLDS: an Ecosystem to Generate, Share and Use Datasets in Reinforcement Learning},
      author={Sabela Ramos and Sertan Girgin and Léonard Hussenot and Damien Vincent and Hanna Yakubovich and Daniel Toyama and Anita Gergely and Piotr Stanczyk and Raphael Marinier and Jeremiah Harmsen and Olivier Pietquin and Nikola Momchev},
      year={2021},
      eprint={2111.02767},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

Acknowledgements

We greatly appreciate all the support from the TF-Agents team in setting up building and testing for EnvLogger.

Disclaimer

This is not an officially supported Google product.