mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-27 14:39:43 +00:00
Merge branch 'main' into feature/add-multitask-dit
This commit is contained in:
@@ -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,
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
Reference in New Issue
Block a user