Add streaming to LeRobotDataset

[Cadene]

What this does

A streaming mode could be useful for training at scale. Instead of downloading all the data locally at one place, data subsets are streamed to training nodes on-demand in a fully distributed fashion. It could also unlock federated learning at a massive scale.

Note: LeRobotDataset has already been used with federated learning tools, but not at a massive scale. See twitter post and flower quickstart lerobot.

This PR adds a streaming: bool argument to LeRobotDataset. Changes are minimal because LeRobotDataset relies on hugging face datasets library and torchvision.io.VideoReader, both already support streaming.

Note: The current implementation of streaming in HF datasets after shuffle is random access.

TODO to unblock merge:

• Compare against webdataset
• Show a real use case where streaming=True is useful
• Add tests

How it was tested

Load single frame from distant video:

from lerobot.common.datasets.video_utils import decode_video_frames_torchvision
video_path = "https://huggingface.co/datasets/pepijn223/mobileso100_drive_forward6/resolve/main/videos/chunk-000/observation.images.mobile/episode_000000.mp4"
frames = decode_video_frames_torchvision(video_path, [0.0], tolerance_s=0.2)

Load items:

from pathlib import Path
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
import tqdm

def main():
    dataset = LeRobotDataset("pepijn223/mobileso100_drive_forward6", episodes=[0], streaming=True)

    for i in tqdm.tqdm(range(100000)):
        item = dataset[i]
        print(f'{item["index"]=}')
        print(f'{item["frame_index"]=}')
        print(f'{item["episode_index"]=}')
        print(f'{item["timestamp"]=}')

if __name__ == "__main__":
    main()

[Tavish9]

Nice feat! Look forward to it 🎊

[aopolin-lv]

Hi, thanks for your effort. However, I notice the streaming mode does not support the 'n_action_step' setting. Do you consider about adding this feature?

[Cadene]

Hi, thanks for your effort. However, I notice the streaming mode does not support the 'n_action_step' setting. Do you consider about adding this feature?

I am sorry I didnt understand. What do you mean? @aopolin-lv