LeRobotDataset v2.1

[aliberts]

What this does

This PR introduces aims to improve the usability of LeRobotDataset. We increase CODEBASE_VERSION from v2.0 to v2.1 as changes are backward compatible with v2.0.

What do I need to do?

Simply run this script on your dataset to update the stats

python lerobot/common/datasets/v21/convert_dataset_v20_to_v21.py \
    --repo-id=repo/id

This will:

• Generate per-episodes stats and writes them in episodes_stats.jsonl
• Check consistency between these new stats and the old ones.
• Remove the deprecated stats.json.
• Update codebase_version in info.json.
• Push this new version to the hub on the main branch and tags it with v2.1.

Changes

• Replaces global stats.json with per-episode stats episodes_stats.jsonl. Episodes stats are then aggregated over selected episodes at initialization of the dataset. Stats computation speed is greatly improved thanks to subsampling of images. Per-episode stats #521

dataset_root/
  ├── data
  ├── meta
  │   ├── episodes.jsonl
+ │   ├── episodes_stats.jsonl
  │   ├── info.json
- │   ├── stats.json
  │   └── tasks.jsonl
  └── videos

• Also adds subsampling for images Add auto downsample height width of images in compute_stats #762
• task is defined per frame and added with add_frame instead of add episode Add frame level task #693
• Ensure simple and consistent typing (np arrays) for the episode_buffer Validate features during add_frame + Add 2D-to-5D + Add string #720
• Remove the need for local_files_only, only use codebase_version from info.json and delete non-main branches Remove local_files_only and use codebase_version instead of branches #734
• Remove the consolidate() method Remove dataset consolidate #752
• Rename LEROBOT_HOME -> HF_LEROBOT_HOME and include HF_HOME Use HF_HOME env variable #753
• Adds option to use HFApi.upload_large_folder for large datasets Add upload_large_folder argument in push_to_hub #763

TODOs in later PRs

• Use standard hf_dataset.set_format("torch") instead of custom hf_dataset.set_transform(hf_transform_to_torch)
• Multi dataset, features mapping
• Update visualization for multi-task episodes

How it was tested

• Improves test_datasets
• Adds test_compute_stats

[Cadene]

Wonderful! Thanks

[Cadene]

Nice catch ;)

[johnMinelli]

I suppose that the following script is not anymore the way to publish dataset+stats ?

lerobot/lerobot/scripts/push_dataset_to_hub.py

Line 223 in ed83cbd

stats = compute_stats(lerobot_dataset, batch_size, num_workers)

[imstevenpmwork]

@johnMinelli Please refer to: #881

[ankile]

Does this change have any implications for the compatibility of the LeRobotDataset and the OnlineBuffer, which is still in verison 2.0 it seems? I.e., if I used LeRobotDataset to pre-train a policy and now I want to continue training using the OnlineBuffer, are there any pitfalls I should be aware of?