LeRobotDataset v2.1 (#711)

Co-authored-by: Remi <remi.cadene@huggingface.co>
Co-authored-by: Remi Cadene <re.cadene@gmail.com>
This commit is contained in:
Simon Alibert
2025-02-25 15:27:29 +01:00
committed by GitHub
parent aca464ca72
commit 3354d919fc
43 changed files with 2023 additions and 1322 deletions

View File

@@ -60,8 +60,6 @@ class RecordControlConfig(ControlConfig):
num_episodes: int = 50
# Encode frames in the dataset into video
video: bool = True
# By default, run the computation of the data statistics at the end of data collection. Compute intensive and not required to just replay an episode.
run_compute_stats: bool = True
# Upload dataset to Hugging Face hub.
push_to_hub: bool = True
# Upload on private repository on the Hugging Face hub.
@@ -83,9 +81,6 @@ class RecordControlConfig(ControlConfig):
play_sounds: bool = True
# Resume recording on an existing dataset.
resume: bool = False
# TODO(rcadene, aliberts): remove local_files_only when refactor with dataset as argument
# Use local files only. By default, this script will try to fetch the dataset from the hub if it exists.
local_files_only: bool = False
def __post_init__(self):
# HACK: We parse again the cli args here to get the pretrained path if there was one.
@@ -130,9 +125,6 @@ class ReplayControlConfig(ControlConfig):
fps: int | None = None
# Use vocal synthesis to read events.
play_sounds: bool = True
# TODO(rcadene, aliberts): remove local_files_only when refactor with dataset as argument
# Use local files only. By default, this script will try to fetch the dataset from the hub if it exists.
local_files_only: bool = False
@ControlConfig.register_subclass("remote_robot")

View File

@@ -183,6 +183,7 @@ def record_episode(
device,
use_amp,
fps,
single_task,
):
control_loop(
robot=robot,
@@ -195,6 +196,7 @@ def record_episode(
use_amp=use_amp,
fps=fps,
teleoperate=policy is None,
single_task=single_task,
)
@@ -210,6 +212,7 @@ def control_loop(
device: torch.device | str | None = None,
use_amp: bool | None = None,
fps: int | None = None,
single_task: str | None = None,
):
# TODO(rcadene): Add option to record logs
if not robot.is_connected:
@@ -224,6 +227,9 @@ def control_loop(
if teleoperate and policy is not None:
raise ValueError("When `teleoperate` is True, `policy` should be None.")
if dataset is not None and single_task is None:
raise ValueError("You need to provide a task as argument in `single_task`.")
if dataset is not None and fps is not None and dataset.fps != fps:
raise ValueError(f"The dataset fps should be equal to requested fps ({dataset['fps']} != {fps}).")
@@ -248,7 +254,7 @@ def control_loop(
action = {"action": action}
if dataset is not None:
frame = {**observation, **action}
frame = {**observation, **action, "task": single_task}
dataset.add_frame(frame)
if display_cameras and not is_headless():