Add video decoding to LeRobotDataset

[Cadene]

What does this PR do?

• Convert datasets into video datasets
• Add video decoding to LeRobotDataset while keeping backward compatibility
• Benchmark available settings and selected: torchvision yuv444p

More information on metrics and settings available here.

Example of benchmark

Effect of compression parameter of ffmpeg video encoding (crf) on:

• compression_factor (higher is better): for instance, compression_factor=4 means that the video takes 4 times less memory space on disk compared to the original images.
• load_time_factor (higher is better): for instance, load_time_factor=0.5 means that decoding from video is 2 times slower than loading the original images.
• avg_per_pixel_l2_error (lower is better): reconstruction error

Note: We selected default value for crf. In fact, as you can see bellow, crf=None allows to compress the image dataset by a factor 15 (15TB -> 1TB), while ensuring acceptable reconstruction error and acceptable loading time reduction (3 times slower).

Note: We also primarily benchmarked -g option which adjusts the amount of key frames. We found -g 2, for 1 key frame every 2 frames, to give the best tradeoff while ensuring high enough loading time:

• 3 times slower for randomly loading 1 frame compared to loading from png images,
• 2 times slower for loading 2 consecutive frames,
• same speed for loading 6 consecutive frames.

Note: In this PR, we use pyav to decode on CPU. We want to explore faster ways to decode like:

• multi-threading,
• GPU decoding.

crf	comp.	load.	avg_l2_error
Baseline	1.0	1.0	0.0000000
0	1.918	0.165	0.0000056
5	3.207	0.171	0.0000111
10	4.818	0.212	0.0000153
15	7.329	0.261	0.0000218
20	11.361	0.312	0.0000317
None	14.932	0.339	0.0000397
25	17.741	0.297	0.0000452
30	27.983	0.406	0.0000629
40	82.449	0.468	0.0001184
50	186.145	0.515	0.0001879

How was it tested?

• proper benchmark on sim and real-world datasets
• reproduced results by retraining from scratch

Pusht

DATA_DIR=data python lerobot/scripts/train.py \
hydra.job.name=pusht_videos

python lerobot/scripts/train.py \
hydra.job.name=pusht_images

Similar success rate, training loss, reward, update time, GPU utilisation.

Aloha

DATA_DIR=data python lerobot/scripts/train.py \
env=aloha \
dataset.repo_id=lerobot/aloha_sim_insertion_human \
policy=act \
hydra.job.name=aloha_sim_insertion_human_videos

python lerobot/scripts/train.py \
env=aloha \
dataset.repo_id=lerobot/aloha_sim_insertion_human \
policy=act \
hydra.job.name=aloha_sim_insertion_human_images

Generate video datasets

python lerobot/scripts/push_dataset_to_hub.py \
--data-dir data \
--dataset-id pusht \
--raw-format pusht_zarr \
--community-id lerobot \
--revision v1.2 \
--dry-run 1 \
--save-to-disk 1 \
--save-tests-to-disk 1

python lerobot/scripts/push_dataset_to_hub.py \
--data-dir data \
--dataset-id xarm_lift_medium \
--raw-format xarm_pkl \
--community-id lerobot \
--revision v1.2 \
--dry-run 1 \
--save-to-disk 1 \
--save-tests-to-disk 1

python lerobot/scripts/push_dataset_to_hub.py \
--data-dir data \
--dataset-id aloha_sim_insertion_human \
--raw-format aloha_hdf5 \
--community-id lerobot \
--revision v1.2 \
--dry-run 1 \
--save-to-disk 1 \
--save-tests-to-disk 1

python lerobot/scripts/push_dataset_to_hub.py \
--data-dir data \
--dataset-id umi_cup_in_the_wild \
--raw-format umi_zarr \
--community-id lerobot \
--revision v1.2 \
--dry-run 1 \
--save-to-disk 1 \
--save-tests-to-disk 1

How to quickly test?

Add --debug 1 to previous push_dataset_to_hub.py commands.

DATA_DIR=tests/data python -m pytest tests/test_datasets.py::test_factory[pusht-lerobot/pusht-diffusion]

TODO

• Video decoding on GPU:
  • decord is used in computer vision papers about videos, but last commit was 2 years ago
  • torchaudio also does GPU decoding through ffmpeg, but I couldnt make it work with yuv444p
  • PyNvCodec is from Nvidia and quite up-to-date. Also they have a nice example with pytorch.

[lhoestq]

Great ! I just left two comments about the video keys. IMO a minor change could be helpful to ensure an easy support for videos in other libs and on HF

[AdilZouitine]

LGTM! This is a very nice PR 📹. I have two minors comment and one suggestion 😄

[alexander-soare]

Looking forward to this landing! Sorry, I think I missed a lot of the action from a previous related PR. Just got nits and comments from me, but generally approved for the time being.

[Cadene]

@alexander-soare @aliberts @AdilZouitine I actually misread the loading time metric. Using the default values of ffmpeg, which are made for movies, leads to extremely slow loading time (25 times slower than images). It makes compute_stats during dataset creation a huge bottleneck. Not acceptable.

I had to redo all benchmarks. I found that by adjusting the number of key frames to be "every 2 frames", we gets a 3 times slower loading time for totally random access of 1 frame, while insuring 15 times reduction in dataset size and acceptable reconstruction error.

Importantly, by retraining, I also validated that this reconstruction error doesnt affect success rate on pusht. I also validated that 3 times slower loading time doesnt affect training time (because we have enough workers).

For TDMPC, loading 6 consecutive frames shouldnt reduce loading time compared to loading png images.
If it is still too slow, we can:

• preload all frames in RAM,
• implement CUDA decoding <--- this require custom torchvision install.

[qgallouedec]

Is it an oversight?