Merge branch 'main' into feature/add-multitask-dit

This commit is contained in:
Bryson Jones
2025-12-15 09:57:27 -08:00
committed by GitHub
2 changed files with 121 additions and 96 deletions
+22 -5
View File
@@ -201,7 +201,8 @@ from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader
from lerobot.utils.control_utils import init_keyboard_listener from lerobot.utils.control_utils import init_keyboard_listener
from lerobot.utils.utils import log_say from lerobot.utils.utils import log_say
from lerobot.utils.visualization_utils import init_rerun from lerobot.utils.visualization_utils import init_rerun
from lerobot.record import record_loop from lerobot.scripts.lerobot_record import record_loop
from lerobot.processor import make_default_processors
NUM_EPISODES = 5 NUM_EPISODES = 5
FPS = 30 FPS = 30
@@ -209,12 +210,19 @@ EPISODE_TIME_SEC = 60
RESET_TIME_SEC = 10 RESET_TIME_SEC = 10
TASK_DESCRIPTION = "My task description" TASK_DESCRIPTION = "My task description"
# Create the robot and teleoperator configurations # Create robot configuration
camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
robot_config = SO100FollowerConfig( robot_config = SO100FollowerConfig(
port="/dev/tty.usbmodem58760434471", id="my_awesome_follower_arm", cameras=camera_config id="my_awesome_follower_arm",
cameras={
"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS) # Optional: fourcc="MJPG" for troubleshooting OpenCV async error.
},
port="/dev/tty.usbmodem58760434471",
)
teleop_config = SO100LeaderConfig(
id="my_awesome_leader_arm",
port="/dev/tty.usbmodem585A0077581",
) )
teleop_config = SO100LeaderConfig(port="/dev/tty.usbmodem585A0077581", id="my_awesome_leader_arm")
# Initialize the robot and teleoperator # Initialize the robot and teleoperator
robot = SO100Follower(robot_config) robot = SO100Follower(robot_config)
@@ -243,6 +251,9 @@ init_rerun(session_name="recording")
robot.connect() robot.connect()
teleop.connect() teleop.connect()
# Create the required processors
teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
episode_idx = 0 episode_idx = 0
while episode_idx < NUM_EPISODES and not events["stop_recording"]: while episode_idx < NUM_EPISODES and not events["stop_recording"]:
log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}") log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
@@ -251,6 +262,9 @@ while episode_idx < NUM_EPISODES and not events["stop_recording"]:
robot=robot, robot=robot,
events=events, events=events,
fps=FPS, fps=FPS,
teleop_action_processor=teleop_action_processor,
robot_action_processor=robot_action_processor,
robot_observation_processor=robot_observation_processor,
teleop=teleop, teleop=teleop,
dataset=dataset, dataset=dataset,
control_time_s=EPISODE_TIME_SEC, control_time_s=EPISODE_TIME_SEC,
@@ -265,6 +279,9 @@ while episode_idx < NUM_EPISODES and not events["stop_recording"]:
robot=robot, robot=robot,
events=events, events=events,
fps=FPS, fps=FPS,
teleop_action_processor=teleop_action_processor,
robot_action_processor=robot_action_processor,
robot_observation_processor=robot_observation_processor,
teleop=teleop, teleop=teleop,
control_time_s=RESET_TIME_SEC, control_time_s=RESET_TIME_SEC,
single_task=TASK_DESCRIPTION, single_task=TASK_DESCRIPTION,
+11 -3
View File
@@ -404,6 +404,10 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
), ),
) )
dataset = None
listener = None
try:
if cfg.resume: if cfg.resume:
dataset = LeRobotDataset( dataset = LeRobotDataset(
cfg.dataset.repo_id, cfg.dataset.repo_id,
@@ -502,14 +506,18 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
dataset.save_episode() dataset.save_episode()
recorded_episodes += 1 recorded_episodes += 1
finally:
log_say("Stop recording", cfg.play_sounds, blocking=True) log_say("Stop recording", cfg.play_sounds, blocking=True)
if dataset:
dataset.finalize()
if robot.is_connected:
robot.disconnect() robot.disconnect()
if teleop is not None: if teleop and teleop.is_connected:
teleop.disconnect() teleop.disconnect()
if not is_headless() and listener is not None: if not is_headless() and listener:
listener.stop() listener.stop()
if cfg.dataset.push_to_hub: if cfg.dataset.push_to_hub: