admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-23 17:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Removed references to batch encoding to be added later or in another PR

Browse Source
This commit is contained in:
Michel Aractingi
2025-07-18 16:52:47 +02:00
parent 5ec70f704e
commit ec40fc41b5
1 changed files with 30 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/record.py
+30 -32
View File
@@ -73,7 +73,6 @@ from lerobot.configs.policies import PreTrainedConfig
from lerobot.datasets.image_writer import safe_stop_image_writer from lerobot.datasets.image_writer import safe_stop_image_writer
from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features
from lerobot.datasets.video_utils import VideoEncodingManager
from lerobot.policies.factory import make_policy from lerobot.policies.factory import make_policy
from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.robots import ( # noqa: F401 from lerobot.robots import ( # noqa: F401
@@ -335,47 +334,46 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
listener, events = init_keyboard_listener() listener, events = init_keyboard_listener()
with VideoEncodingManager(dataset): recorded_episodes = 0
recorded_episodes = 0 while recorded_episodes < cfg.dataset.num_episodes and not events["stop_recording"]:
while recorded_episodes < cfg.dataset.num_episodes and not events["stop_recording"]: log_say(f"Recording episode {dataset.num_episodes}", cfg.play_sounds)
log_say(f"Recording episode {dataset.num_episodes}", cfg.play_sounds) record_loop(
robot=robot,
events=events,
fps=cfg.dataset.fps,
teleop=teleop,
policy=policy,
dataset=dataset,
control_time_s=cfg.dataset.episode_time_s,
single_task=cfg.dataset.single_task,
display_data=cfg.display_data,
)
# Execute a few seconds without recording to give time to manually reset the environment
# Skip reset for the last episode to be recorded
if not events["stop_recording"] and (
(recorded_episodes < cfg.dataset.num_episodes - 1) or events["rerecord_episode"]
):
log_say("Reset the environment", cfg.play_sounds)
record_loop( record_loop(
robot=robot, robot=robot,
events=events, events=events,
fps=cfg.dataset.fps, fps=cfg.dataset.fps,
teleop=teleop, teleop=teleop,
policy=policy, control_time_s=cfg.dataset.reset_time_s,
dataset=dataset,
control_time_s=cfg.dataset.episode_time_s,
single_task=cfg.dataset.single_task, single_task=cfg.dataset.single_task,
display_data=cfg.display_data, display_data=cfg.display_data,
) )
# Execute a few seconds without recording to give time to manually reset the environment if events["rerecord_episode"]:
# Skip reset for the last episode to be recorded log_say("Re-record episode", cfg.play_sounds)
if not events["stop_recording"] and ( events["rerecord_episode"] = False
(recorded_episodes < cfg.dataset.num_episodes - 1) or events["rerecord_episode"] events["exit_early"] = False
): dataset.clear_episode_buffer()
log_say("Reset the environment", cfg.play_sounds) continue
record_loop(
robot=robot,
events=events,
fps=cfg.dataset.fps,
teleop=teleop,
control_time_s=cfg.dataset.reset_time_s,
single_task=cfg.dataset.single_task,
display_data=cfg.display_data,
)
if events["rerecord_episode"]: dataset.save_episode()
log_say("Re-record episode", cfg.play_sounds) recorded_episodes += 1
events["rerecord_episode"] = False
events["exit_early"] = False
dataset.clear_episode_buffer()
continue
dataset.save_episode()
recorded_episodes += 1
log_say("Stop recording", cfg.play_sounds, blocking=True) log_say("Stop recording", cfg.play_sounds, blocking=True)