mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-22 03:59:42 +00:00
make lerobot record easier
This commit is contained in:
@@ -487,18 +487,9 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
|
|||||||
robot = make_robot_from_config(cfg.robot)
|
robot = make_robot_from_config(cfg.robot)
|
||||||
teleop = make_teleoperator_from_config(cfg.teleop) if cfg.teleop is not None else None
|
teleop = make_teleoperator_from_config(cfg.teleop) if cfg.teleop is not None else None
|
||||||
|
|
||||||
# Dataset features derived automatically from robot/teleop pipelines
|
# Dataset features derived automatically from robot/teleop pipelines.
|
||||||
if teleop is not None:
|
# When teleop is None (policy-only recording), only observation features are included.
|
||||||
dataset_features = build_dataset_features(robot, teleop, use_videos=cfg.dataset.video)
|
dataset_features = build_dataset_features(robot, teleop, use_videos=cfg.dataset.video)
|
||||||
else:
|
|
||||||
# Policy-only recording: use robot observation features only
|
|
||||||
from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features, create_initial_features
|
|
||||||
|
|
||||||
dataset_features = aggregate_pipeline_dataset_features(
|
|
||||||
pipeline=robot.output_pipeline(),
|
|
||||||
initial_features=create_initial_features(observation=robot.raw_observation_features),
|
|
||||||
use_videos=cfg.dataset.video,
|
|
||||||
)
|
|
||||||
|
|
||||||
dataset = None
|
dataset = None
|
||||||
listener = None
|
listener = None
|
||||||
|
|||||||
@@ -27,7 +27,7 @@ from lerobot.datasets.utils import combine_feature_dicts
|
|||||||
|
|
||||||
def build_dataset_features(
|
def build_dataset_features(
|
||||||
robot,
|
robot,
|
||||||
teleop,
|
teleop=None,
|
||||||
*,
|
*,
|
||||||
use_videos: bool = True,
|
use_videos: bool = True,
|
||||||
) -> dict:
|
) -> dict:
|
||||||
@@ -35,15 +35,16 @@ def build_dataset_features(
|
|||||||
Derive dataset feature specifications from robot and teleoperator pipelines.
|
Derive dataset feature specifications from robot and teleoperator pipelines.
|
||||||
|
|
||||||
Uses the robot's ``output_pipeline`` and ``raw_observation_features`` to determine
|
Uses the robot's ``output_pipeline`` and ``raw_observation_features`` to determine
|
||||||
what the dataset will store as observations, and the teleoperator's ``output_pipeline``
|
what the dataset will store as observations, and (when provided) the teleoperator's
|
||||||
and ``raw_action_features`` to determine what will be stored as actions.
|
``output_pipeline`` and ``raw_action_features`` to determine what will be stored as actions.
|
||||||
|
|
||||||
This replaces the old pattern of manually calling ``aggregate_pipeline_dataset_features``
|
This replaces the old pattern of manually calling ``aggregate_pipeline_dataset_features``
|
||||||
with explicit processor objects.
|
with explicit processor objects.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
robot: The robot instance (must have ``output_pipeline()`` and ``raw_observation_features``).
|
robot: The robot instance (must have ``output_pipeline()`` and ``raw_observation_features``).
|
||||||
teleop: The teleoperator instance (must have ``output_pipeline()`` and ``raw_action_features``).
|
teleop: The teleoperator instance. When ``None`` (policy-only recording), only observation
|
||||||
|
features are returned.
|
||||||
use_videos: If True, image observations are included as video features.
|
use_videos: If True, image observations are included as video features.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
@@ -51,17 +52,19 @@ def build_dataset_features(
|
|||||||
|
|
||||||
Example::
|
Example::
|
||||||
|
|
||||||
dataset = LeRobotDataset.create(
|
# Teleop recording
|
||||||
repo_id="...",
|
features = build_dataset_features(follower, leader, use_videos=True)
|
||||||
fps=30,
|
|
||||||
features=build_dataset_features(follower, leader, use_videos=True),
|
# Policy-only recording (no teleop)
|
||||||
)
|
features = build_dataset_features(robot, use_videos=True)
|
||||||
"""
|
"""
|
||||||
obs_features = aggregate_pipeline_dataset_features(
|
obs_features = aggregate_pipeline_dataset_features(
|
||||||
pipeline=robot.output_pipeline(),
|
pipeline=robot.output_pipeline(),
|
||||||
initial_features=create_initial_features(observation=robot.raw_observation_features),
|
initial_features=create_initial_features(observation=robot.raw_observation_features),
|
||||||
use_videos=use_videos,
|
use_videos=use_videos,
|
||||||
)
|
)
|
||||||
|
if teleop is None:
|
||||||
|
return obs_features
|
||||||
action_features = aggregate_pipeline_dataset_features(
|
action_features = aggregate_pipeline_dataset_features(
|
||||||
pipeline=teleop.output_pipeline(),
|
pipeline=teleop.output_pipeline(),
|
||||||
initial_features=create_initial_features(action=teleop.raw_action_features),
|
initial_features=create_initial_features(action=teleop.raw_action_features),
|
||||||
|
|||||||
Reference in New Issue
Block a user