RLDS (Reinforcement Learning Datasets)

RLDS (Reinforcement Learning Datasets) is an ecosystem of tools designed to store, retrieve, and manipulate episodic data for Sequential Decision Making tasks, such as Reinforcement Learning (RL), Offline RL, and Imitation Learning.

Key components of the ecosystem include:

• RLDS Library: For manipulating RLDS-compliant datasets.
• EnvLogger: For creating synthetic datasets.
• RLDS Creator: For creating datasets from human-environment interactions.
• TFDS (TensorFlow Datasets): For accessing existing RL datasets and loading RLDS datasets.

Even if your data is already in TFDS format, you may want to implement a formal TFDS builder for two main reasons:

• Reshuffling: To re-generate the data so that episodes are shuffled on disk. By default, episodes are stored in the order they were generated by Envlogger.
• Sharing: To either add the dataset to the official TFDS catalog or to share it within your own repository. Users can still load data directly using tfds.builder_from_directory without a formal builder.

When using multiple processes to read the same dataset (e.g., multiple actors in a single-learner scenario), you must ensure processes do not receive the same sequence of episodes. There are two primary strategies:

1. Using TFDS Splits (Deterministic): Use the split API to assign a unique, disjoint set of episodes to each reader. This is the easiest method, but if a reader process dies, its assigned portion of the dataset will not be processed.
2. Non-deterministic Reading: Set shuffle_files=True and tune ReadConfig options in tfds.load or builder.as_dataset. This allows the full dataset to be processed even if a reader dies, but some episodes may appear more than once.

If your algorithm operates on individual steps or n-step transitions rather than full episodes, you can interleave steps from multiple episodes to achieve randomization.

One method is to use tf.data.Dataset.interleave to create N copies of the dataset. Each copy shuffles input partitions independently, ensuring consecutive steps come from unrelated episodes.

Note: To avoid loading the same step N times, construct each copy using disjoint TFDS splits.

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=...)

For practical implementations and common patterns in Reinforcement Learning, use the RL examples Colab.

https://colab.research.google.com/github/google-research/rlds/blob/main/rlds/examples/rlds_examples.ipynb

Since RLDS datasets are exposed as tf.data.Dataset objects, standard TensorFlow performance optimizations apply. However, because RLDS datasets have a specific structure, some general speed-up methods may not work as expected.

To improve performance:

1. Use the optimized RLDS library of transformations.
2. Follow the best practices outlined in the performance Colab.

If you are experiencing bottlenecks or want to optimize your data loading pipeline, refer to the Performance Best Practices Colab. It provides guidance on how to use RLDS efficiently.

https://colab.research.google.com/github/google-research/rlds/blob/main/rlds/examples/rlds_performance.ipynb

If your data is in a custom format that is neither RLDS nor TFDS compatible, follow these two steps to integrate it into TFDS:

1. Implement a Python dataset builder class: Define the dataset specifications (e.g., shapes of observations, actions, etc.) and the logic for reading your raw data files.
   • Note: To ensure compatibility with RLDS pipelines, your implementation must provide the same structure and keys as an RLDS dataset.
2. Run the download_and_prepare pipeline: This converts your raw data into the TFDS intermediate format.

Alternatively, you can rewrite raw data into an RLDS/TFDS compatible format before adding it to TFDS using Envlogger or EpisodeWriter. For EpisodeWriter, use the ConfigGenerator tool to create your DatasetConfig.