TFDS (TensorFlow Datasets): Complete Guide for Robotics Data

TFDS (TensorFlow Datasets) extends TFRecord with a standardized builder pattern that combines data serialization, schema definition, versioning, and pipeline integration into a single framework. A DatasetBuilder subclass defines a DatasetInfo object containing the feature structure (tfds.features.Image for images with automatic JPEG/PNG encoding, tfds.features.Tensor for numerical arrays with specified dtype and shape, tfds.features.Text for strings, tfds.features.FeaturesDict for nested dictionaries), the semantic version number (e.g., 1.0.0), and a human-readable description. Data is stored as sharded TFRecords (typically 256 MB per shard) with a dataset_info.json manifest that records the feature schema, split statistics, number of examples, and file checksums.

The TFDS builder pattern provides a standardized way to define, generate, and version datasets. A DatasetBuilder subclass implements three key methods: _info() returns the DatasetInfo with feature connectors and metadata, _split_generators() defines how to discover and partition raw data files into train/validation/test splits, and _generate_examples() yields individual examples as (key, feature_dict) pairs that the framework serializes to sharded TFRecords with automatic checksumming. The builder system handles shard balancing (distributing examples evenly across shards), deterministic shuffling (consistent ordering across builds), and download management (caching raw data and tracking URLs). Version numbers follow semantic versioning: major version changes indicate backward-incompatible feature changes, minor versions add features, and patch versions fix data errors.

TFDS's integration with tf.data provides production-grade data pipeline primitives that are critical for efficient GPU utilization during training. The tfds.load() function returns a tf.data.Dataset with configurable batching, shuffling (with configurable buffer size), prefetching (overlapping data loading with GPU computation), interleaving (reading from multiple shards simultaneously), and caching (keeping frequently accessed data in memory). For robotics pipelines, the combination of RLDS (which defines the episode/step structure on top of TFDS) with TFDS (which handles the storage and pipeline mechanics) has become the standard in the Google DeepMind ecosystem, powering RT-1, RT-2, Octo, and OpenVLA training runs. TFDS also supports loading datasets from Google Cloud Storage with automatic local caching, enabling efficient access to multi-terabyte datasets without full downloads.

Loading a TFDS dataset requires just two lines: dataset = tfds.load('dataset_name', split='train') returns a tf.data.Dataset, and iterating with for example in dataset: accesses individual examples as nested dictionaries of tensors. For robotics RLDS datasets: dataset = tfds.load('bridge_dataset', split='train') returns the full BridgeData V2 collection. Each example is an episode containing a 'steps' dataset that you iterate: for episode in dataset: for step in episode['steps']: where step['observation']['image'] yields a decoded image tensor and step['action'] yields the action vector. The tfds.load() function handles downloading, caching, shuffling, and batching automatically.

Creating a new TFDS dataset involves writing a DatasetBuilder subclass. The _info method returns tfds.core.DatasetInfo with a features dictionary: tfds.features.FeaturesDict({'image': tfds.features.Image(shape=(H, W, 3)), 'action': tfds.features.Tensor(shape=(7,), dtype=tf.float32), 'reward': tf.float32}). The _split_generators method points to raw data directories, and _generate_examples yields (example_key, example_dict) tuples. Running tfds build --datasets=your_dataset generates the sharded TFRecords, dataset_info.json, and label metadata. For the Open X-Embodiment ecosystem, builders follow the RLDS convention of yielding episodes as nested datasets of steps, and the repository at github.com/google-deepmind/open_x_embodiment maintains reference builder implementations.

TFDS provides several advanced features critical for large-scale robotics training. The tfds.ReadConfig class controls reading behavior: num_parallel_reads (default 64) determines how many shards are read concurrently, interleave_cycle_length controls round-robin shard reading for better shuffling, and try_autocache enables automatic in-memory caching for datasets that fit in RAM. For distributed training across multiple TPU/GPU hosts, TFDS integrates with tf.distribute strategies to automatically shard data across workers, ensuring each worker processes a unique subset of examples. The tfds.benchmarks module provides throughput measurement tools, and well-configured TFDS pipelines achieve 1-10 GB/s data throughput on modern storage systems.

TFDS is the foundational data layer for Google DeepMind's robotics research stack. The Open X-Embodiment project, which aggregates over 60 robotics datasets from 21 institutions into a unified collection for training cross-embodiment policies, uses TFDS builders for every dataset. Each contributing lab writes a DatasetBuilder that converts their native format (HDF5, custom CSVs, ROS bags) into a standardized RLDS-on-TFDS representation. This standardization enables models like RT-X and Octo to train on the entire collection with a single data loading pipeline, despite the extreme heterogeneity of the underlying data sources.

The combination of TFDS and tf.data provides critical performance optimizations for large-scale robotics model training. TFDS's deterministic shuffling ensures reproducible training runs (important for ablation studies), while tf.data's prefetching and interleaving pipelines keep GPU utilization above 90% even when reading from networked storage. For TPU training (used by RT-2 and RT-X), TFDS's integration with tf.data.service enables distributed data preprocessing where CPU-bound operations (image decoding, augmentation) run on separate preprocessing workers, freeing TPU host CPUs for feeding data to accelerators.

For teams outside the Google ecosystem who want to use TFDS datasets, several interoperability paths exist. The tfds.as_numpy() function strips TensorFlow tensors to NumPy arrays suitable for any framework. The dlimp library (used by the Octo team) wraps TFDS datasets in PyTorch-compatible iterators. The fog_x library provides a format-agnostic abstraction layer that can read both TFDS and other robotics formats. For teams that want to consume Open X-Embodiment data in PyTorch, Claru can convert any RLDS/TFDS dataset to LeRobot format (Parquet + MP4) for native HuggingFace Hub compatibility.

Claru provides custom TFDS DatasetBuilder code alongside the data, enabling one-line loading via tfds.load('claru_your_dataset'). Builders include properly typed feature connectors (tfds.features.Image with configurable encoding, tfds.features.Tensor with exact dtype and shape specifications), split configurations matching your train/val/test requirements, and comprehensive dataset documentation embedded in the DatasetInfo. Image features use JPEG compression by default with configurable quality (default 95), and proprioceptive state uses float32 precision.

Every delivery includes the complete builder Python package that can be registered in your local TFDS installation, a set of pre-built TFRecord shards with checksums, and a verification script that runs the builder's built-in integrity checks. For teams targeting the Open X-Embodiment ecosystem, we validate schema compatibility with RT-X and Octo model expectations (observation dictionary structure, action space dimensionality, language instruction format). For teams needing PyTorch compatibility, we additionally deliver a dlimp-compatible wrapper and can provide the same data in LeRobot or HDF5 format.