From 651e9ca313bda76cee43054ef3de1cff24d7c340 Mon Sep 17 00:00:00 2001 From: Saberlve Date: Fri, 7 Aug 2026 14:14:55 +0800 Subject: [PATCH] Add manual mode data collection script --- .gitignore | 1 + README.md | 17 +- README_ZH.md | 17 +- .../xarm7_manual_record_config.yaml | 32 +++ pyproject.toml | 1 + .../robots/uf_robot/uf_robot.py | 22 ++ .../robots/uf_robot/uf_robot_config.py | 7 + .../scripts/uf_lerobot_record.py | 100 +++++++-- tests/test_manual_mode.py | 203 ++++++++++++++++++ 9 files changed, 378 insertions(+), 22 deletions(-) create mode 100644 config/manual_mode/xarm7_manual_record_config.yaml create mode 100644 tests/test_manual_mode.py diff --git a/.gitignore b/.gitignore index 10d90fb..1f82ff6 100644 --- a/.gitignore +++ b/.gitignore @@ -88,3 +88,4 @@ models/ *.trt *.xvcd ufactory_usage/ +.history/ \ No newline at end of file diff --git a/README.md b/README.md index ff9f6c3..5fdb4e1 100644 --- a/README.md +++ b/README.md @@ -165,7 +165,18 @@ uv run uf-xarm-manual-mode --config_path config/manual_mode/xarm_manual_mode_con Set `manual_mode` to `true` to enter drag mode. Press Enter to restore normal mode; set it to `false` to restore normal mode directly. With `return_to_initial: true`, the script reads the saved initial point through the xArm Studio API and returns to it when drag mode exits. -### 3. Data Collection +### 3. Manual Drag Data Collection + +Manual drag recording uses `manual_mode: true` in the robot configuration and does not configure a teleoperator. During recording, the actual joint state is written as both the observation and action in the LeRobot dataset: + +```bash +uv run uf-lerobot-record --config_path config/manual_mode/xarm7_manual_record_config.yaml +uv run uf-lerobot-record --config_path config/manual_mode/xarm7_manual_record_config.yaml --resume true +``` + +Press `Space` to start an episode, `Right` to save it, `Left` to discard and re-record it, and `Esc` to stop recording. Reset the arm manually between episodes. + +### 4. Teleop Data Collection Record datasets via teleop. @@ -178,7 +189,7 @@ uv run uf-lerobot-record --config_path path/to/config.yaml --resume true uv run uf-lerobot-record --config_path config/umi/xarm6_umi_record_config.yaml ``` -### 4. Policy Training +### 5. Policy Training Train imitation learning policies on collected data. @@ -209,7 +220,7 @@ uv run lerobot-train \ --save_freq=20000 ``` -### 5. Inference & Evaluation +### 6. Inference & Evaluation Run inference with a trained policy. diff --git a/README_ZH.md b/README_ZH.md index 7adcd4d..95f0de4 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -164,7 +164,18 @@ uv run uf-xarm-manual-mode --config_path config/manual_mode/xarm_manual_mode_con 配置中的 `manual_mode` 为 `true` 时进入拖拽模式，按回车后恢复为普通模式；设置为 `false` 可直接恢复为普通模式。`return_to_initial: true` 会通过 xArm Studio 接口读取保存的初始位置，并在退出拖拽后自动回位。 -### 3. 数据采集 +### 3. 人工拖拽数据采集 + +人工拖拽录制使用 robot 配置中的 `manual_mode: true`，不需要配置 teleop。录制过程中，机械臂的实际关节状态会同时作为 observation 和 action 写入 LeRobot 数据集： + +```bash +uv run uf-lerobot-record --config_path config/manual_mode/xarm7_manual_record_config.yaml +uv run uf-lerobot-record --config_path config/manual_mode/xarm7_manual_record_config.yaml --resume true +``` + +按 `Space` 开始当前 episode，按 `Right` 保存，按 `Left` 放弃并重录当前 episode，按 `Esc` 停止录制。episode 之间可以手动复位机械臂。 + +### 4. 遥操作数据采集 通过遥操作录制数据集。 @@ -177,7 +188,7 @@ uv run uf-lerobot-record --config_path path/to/config.yaml --resume true # uv run uf-lerobot-record --config_path config/umi/xarm6_umi_record_config.yaml ``` -### 4. Lerobot训练 +### 5. Lerobot训练 采集数据后，使用 LeRobot 训练管道进行模仿学习训练。 @@ -205,7 +216,7 @@ uv run lerobot-train \ --save_freq=20000 ``` -### 5. 推理 +### 6. 推理 指定模型进行推理 diff --git a/config/manual_mode/xarm7_manual_record_config.yaml b/config/manual_mode/xarm7_manual_record_config.yaml new file mode 100644 index 0000000..6abf68d --- /dev/null +++ b/config/manual_mode/xarm7_manual_record_config.yaml @@ -0,0 +1,32 @@ +robot: + type: uf::robot + id: "uf_manual_robot" + robot_dof: 7 + control_space: "joint" + # IP address of the xArm controller. + robot_ip: "192.168.1.245" + # Gripper type: 1 is the xArm gripper. + gripper_type: 1 + manual_mode: true + # Teaching sensitivity, valid range is 1-5, affecting vel of the robot. + teach_sensitivity: 3 + # Whether to record joint velocities in observations. + observe_joint_vel: false + + # Camera definitions; leave empty when no cameras are used. + cameras: {} + +dataset: + root: "/home/uf/Data/lerobot_datas/record/ufactory/xarm7_manual_datas" + repo_id: "ufactory/xarm7_manual_datas" + # Task description stored with each recorded frame. + single_task: "Describe the task being demonstrated." + fps: 30 + episode_time_s: 60 + # Reserved reset duration between episodes in seconds. + reset_time_s: 20 + # Number of episodes to record. + num_episodes: 50 + # Store camera observations as videos. + video: true + push_to_hub: false diff --git a/pyproject.toml b/pyproject.toml index 3e2d914..2713a72 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,7 @@ uf-xarm-manual-mode = "lerobot_robot_ufactory.scripts.uf_xarm_manual_mode:main" uf-lerobot-eval = "lerobot_robot_ufactory.scripts.uf_lerobot_eval:main" uf-vive-calibrate = "lerobot_robot_ufactory.scripts.vive_calibrate:main" uf-camera-view = "lerobot_robot_ufactory.scripts.uf_camera_view:main" +uf-camera-test = "lerobot_robot_ufactory.scripts.uf_camera_test:main" [project.optional-dependencies] # GELLO 遥操作 diff --git a/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot.py b/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot.py index db7fd67..d50efcc 100644 --- a/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot.py +++ b/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot.py @@ -221,6 +221,26 @@ class UFRobot(Robot, Thread): self.real_arm.set_state(0) # set to start state time.sleep(0.5) + if self.config.manual_mode: + if self.config.teach_sensitivity is not None: + code = self.real_arm.set_teach_sensitivity(self.config.teach_sensitivity) + if code != 0: + raise RuntimeError(f"set_teach_sensitivity failed, code={code}") + + code = self.real_arm.set_mode(2) + if code != 0: + raise RuntimeError(f"set_mode(2) failed, code={code}") + code = self.real_arm.set_state(0) + if code != 0: + raise RuntimeError(f"set_state(0) failed, code={code}") + + _, err_warn = self.real_arm.get_err_warn_code() + if err_warn[0] != 0: + raise RuntimeError( + f"Failed to set manual mode for UF robot! Controller Error code: {err_warn[0]} !" + ) + return + _, err_warn = self.real_arm.get_err_warn_code() if err_warn[0] != 0: raise RuntimeError(f"Failed to set correct state to UF robot! Controller Error code: {err_warn[0]} !") @@ -353,6 +373,8 @@ class UFRobot(Robot, Thread): def send_action(self, action: dict) -> np.ndarray: if not self._is_connected: raise ConnectionError() + if self.config.manual_mode: + return action if self.real_arm.error_code != 0: return action if self.config.no_action: diff --git a/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot_config.py b/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot_config.py index 590c11c..f170f4b 100644 --- a/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot_config.py +++ b/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot_config.py @@ -18,6 +18,8 @@ class UFRobotConfig(RobotConfig): gripper_speed: int = -1 # auto gripper_force: int = -1 # auto observe_joint_vel: bool = False # only effective in joint control mode + manual_mode: bool = False # xArm joint teaching mode; records state without sending actions + teach_sensitivity: int | None = None # xArm teaching sensitivity, valid range: 1-5 start_joints: Tuple[float, ...] = (0, 0, 0, 90, 0, 90, 0) # ° start_tcp_pose: Tuple[float, ...] = None # [x, y, z, roll(°), pitch(°), yaw(°)] max_joint_velocity: int = 90 # °/s, only effective in joint control mode @@ -27,3 +29,8 @@ class UFRobotConfig(RobotConfig): def __post_init__(self): super().__post_init__() self.id = 'uf_robot' if self.id is None else self.id + if self.manual_mode: + if self.control_space != "joint": + raise ValueError("manual_mode requires control_space='joint'") + if self.teach_sensitivity is not None and not 1 <= self.teach_sensitivity <= 5: + raise ValueError("teach_sensitivity must be between 1 and 5") diff --git a/src/lerobot_robot_ufactory/scripts/uf_lerobot_record.py b/src/lerobot_robot_ufactory/scripts/uf_lerobot_record.py index 6549d7e..f458d59 100644 --- a/src/lerobot_robot_ufactory/scripts/uf_lerobot_record.py +++ b/src/lerobot_robot_ufactory/scripts/uf_lerobot_record.py @@ -6,14 +6,29 @@ import argparse import logging import shutil import threading +from dataclasses import dataclass from pathlib import Path import lerobot_robot_ufactory # patch from lerobot.scripts.lerobot_record import * +from lerobot.scripts.lerobot_record import RecordConfig as LeRobotRecordConfig from lerobot_robot_ufactory.teleoperators.uf_mock_teleop import UFMockTeleop from lerobot_robot_ufactory.teleoperators.base_teleop import UFBaseTeleop from lerobot_robot_ufactory.utils.utils import init_keyboard_listener +@dataclass +class UFRecordConfig(LeRobotRecordConfig): + """RecordConfig variant that permits UFACTORY manual-mode recording.""" + + def __post_init__(self): + manual_mode = getattr(self.robot, "manual_mode", False) + if manual_mode: + if self.teleop is not None or self.policy is not None: + raise ValueError("manual_mode recording cannot be combined with a teleop or policy") + return + super().__post_init__() + + def _get_dataset_writer(dataset): return getattr(dataset, "writer", None) @@ -81,6 +96,11 @@ def _current_episode_index(dataset): return dataset.num_episodes +def _manual_action_from_observation(observation, action_features): + """Keep only robot action fields when mirroring manual-mode state.""" + return {key: value for key, value in observation.items() if key in action_features} + + def _create_empty_episode_buffer(dataset, episode_index, template_episode_buffer): writer = _get_dataset_writer(dataset) @@ -121,6 +141,7 @@ class AsyncEpisodeSaver: self._total_cnts = 0 self._finish_cnts = 0 self._exception = None + self._closed = False self._thread = threading.Thread(target=self._run, name="uf-async-episode-saver", daemon=True) self._thread.start() @@ -141,10 +162,13 @@ class AsyncEpisodeSaver: self._raise_if_failed() def close(self): + if self._closed: + return self._queue.join() self._queue.put(self._STOP) self._queue.join() self._thread.join() + self._closed = True self._raise_if_failed() def _run(self): @@ -187,6 +211,39 @@ class AsyncEpisodeSaver: def _raise_if_failed(self): if self._exception is not None: raise RuntimeError("Async episode save failed.") from self._exception + + +def _disconnect_recording_resources(robot, teleop, listener): + """Release recording devices while preserving cleanup after partial failures.""" + try: + if getattr(robot, "_is_connected", False) or getattr(robot, "real_arm", None) is not None: + robot.disconnect() + finally: + try: + if teleop is not None and getattr(teleop, "is_connected", False): + teleop.disconnect() + finally: + if listener is not None: + listener.stop() + + +class _RecordingCleanup: + def __init__(self, robot, teleop, listener, async_episode_saver): + self.robot = robot + self.teleop = teleop + self.listener = listener + self.async_episode_saver = async_episode_saver + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + try: + if self.async_episode_saver is not None: + self.async_episode_saver.close() + finally: + _disconnect_recording_resources(self.robot, self.teleop, self.listener) + return False @safe_stop_image_writer @@ -213,6 +270,7 @@ def record_loop( display_data: bool = False, display_compressed_images: bool = False, frame_callback: callable = None, + manual_mode: bool = False, ): if dataset is not None and dataset.fps != fps: raise ValueError(f"The dataset fps should be equal to requested fps ({dataset.fps} != {fps}).") @@ -285,6 +343,12 @@ def record_loop( act_processed_policy: RobotAction = make_robot_action(action_values, dataset.features) + elif policy is None and manual_mode: + # In manual mode the physical arm is the source of both the + # observation and the demonstrated target state. + act = _manual_action_from_observation(obs_processed, robot.action_features) + act_processed_teleop = teleop_action_processor((act, obs)) + elif policy is None and isinstance(teleop, Teleoperator): act = teleop.get_action() @@ -344,7 +408,7 @@ def record_loop( timestamp = time.perf_counter() - start_episode_t -def record(cfg: RecordConfig, async_save: bool = False) -> LeRobotDataset: +def record(cfg: UFRecordConfig, async_save: bool = False) -> LeRobotDataset: init_logging() logging.info(pformat(asdict(cfg))) if cfg.display_data: @@ -352,6 +416,7 @@ def record(cfg: RecordConfig, async_save: bool = False) -> LeRobotDataset: robot = make_robot_from_config(cfg.robot) teleop = make_teleoperator_from_config(cfg.teleop) if cfg.teleop is not None else None + manual_mode = bool(getattr(cfg.robot, "manual_mode", False)) teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors() @@ -413,14 +478,22 @@ def record(cfg: RecordConfig, async_save: bool = False) -> LeRobotDataset: }, ) - robot.connect() - if teleop is not None: - teleop.connect() + try: + robot.connect() + if teleop is not None: + teleop.connect() + except BaseException: + try: + _disconnect_recording_resources(robot, teleop, None) + except BaseException: + logging.exception("Failed to clean up after recording device connection failure") + raise is_evt = not is_headless() is_uf_teleop = isinstance(teleop, UFBaseTeleop) is_recorded = False key_dict = {} + listener = None events = {"exit_early": False, "rerecord_episode": False, "stop_recording": False} if is_evt: @@ -477,7 +550,7 @@ def record(cfg: RecordConfig, async_save: bool = False) -> LeRobotDataset: if async_episode_saver is not None: print('Async episode saving is enabled.') - with VideoEncodingManager(dataset): + with _RecordingCleanup(robot, teleop, listener, async_episode_saver), VideoEncodingManager(dataset): recorded_episodes = 0 while recorded_episodes < cfg.dataset.num_episodes and not events["stop_recording"]: time.sleep(0.01) @@ -501,10 +574,11 @@ def record(cfg: RecordConfig, async_save: bool = False) -> LeRobotDataset: if is_recorded: events["rerecord_episode"] = False events["exit_early"] = False - if is_uf_teleop: + if is_uf_teleop or manual_mode: robot.configure() - obs = robot.get_observation() - teleop.set_teleop_enabled(True, obs) + if is_uf_teleop: + obs = robot.get_observation() + teleop.set_teleop_enabled(True, obs) log_say(f"Recording episode {_current_episode_index(dataset)}", cfg.play_sounds) record_loop( robot=robot, @@ -522,6 +596,7 @@ def record(cfg: RecordConfig, async_save: bool = False) -> LeRobotDataset: single_task=cfg.dataset.single_task, display_data=cfg.display_data, frame_callback=frame_callback, + manual_mode=manual_mode, ) else: continue @@ -576,13 +651,6 @@ def record(cfg: RecordConfig, async_save: bool = False) -> LeRobotDataset: print("\n********** Episode Record Loop Exit **********") - robot.disconnect() - if teleop is not None: - teleop.disconnect() - - if is_evt and listener is not None: - listener.stop() - if cfg.dataset.push_to_hub: dataset.push_to_hub(tags=cfg.dataset.tags, private=cfg.dataset.private) @@ -590,7 +658,7 @@ def record(cfg: RecordConfig, async_save: bool = False) -> LeRobotDataset: return dataset @parser.wrap() -def get_cfg(cfg: RecordConfig) -> RecordConfig: +def get_cfg(cfg: UFRecordConfig) -> UFRecordConfig: return cfg def main(): diff --git a/tests/test_manual_mode.py b/tests/test_manual_mode.py new file mode 100644 index 0000000..6a57b96 --- /dev/null +++ b/tests/test_manual_mode.py @@ -0,0 +1,203 @@ +import sys +from pathlib import Path + +import numpy as np +import pytest +from lerobot.datasets.lerobot_dataset import LeRobotDataset + +from lerobot_robot_ufactory.robots.uf_robot.uf_robot_config import UFRobotConfig +from lerobot_robot_ufactory.scripts import uf_lerobot_record as record_module +from lerobot_robot_ufactory.scripts.uf_lerobot_record import ( + _manual_action_from_observation, + get_cfg, +) + + +class FakeXArm: + def __init__(self, robot_ip): + self.robot_ip = robot_ip + self.connected = True + self.axis = 6 + self.error_code = 0 + self.mode = 0 + self.calls = [] + + def motion_enable(self, **kwargs): + self.calls.append(("motion_enable", kwargs)) + + def clean_error(self): + self.calls.append(("clean_error",)) + + def set_teach_sensitivity(self, sensitivity): + self.calls.append(("set_teach_sensitivity", sensitivity)) + return 0 + + def set_mode(self, mode): + self.calls.append(("set_mode", mode)) + self.mode = mode + return 0 + + def set_state(self, state): + self.calls.append(("set_state", state)) + return 0 + + def get_err_warn_code(self): + return 0, [0, 0] + + def set_linear_spd_limit_factor(self, factor): + self.calls.append(("set_linear_spd_limit_factor", factor)) + return 0 + + def get_joint_states(self, is_radian=True, num=3): + positions = np.arange(6, dtype=np.float64) + velocities = np.zeros(6, dtype=np.float64) + return 0, [positions, velocities, velocities] + + def disconnect(self): + self.calls.append(("disconnect",)) + self.connected = False + + +def test_manual_mode_robot_enters_teaching_mode_without_sending_actions(monkeypatch, tmp_path): + from lerobot_robot_ufactory.robots.uf_robot import uf_robot as uf_robot_module + + arm = FakeXArm("192.168.1.245") + monkeypatch.setattr(uf_robot_module, "XArmAPI", lambda robot_ip: arm) + monkeypatch.setattr(uf_robot_module.time, "sleep", lambda _: None) + + config = UFRobotConfig( + id="test_manual_robot", + calibration_dir=tmp_path, + robot_ip=arm.robot_ip, + robot_dof=6, + control_space="joint", + gripper_type=0, + start_joints=(), + manual_mode=True, + teach_sensitivity=4, + ) + robot = uf_robot_module.UFRobot(config) + + robot.connect() + assert arm.mode == 2 + assert ("set_teach_sensitivity", 4) in arm.calls + + action = {"J1.pos": 1.0} + assert robot.send_action(action) is action + assert not any(call[0] == "set_servo_angle" for call in arm.calls) + + observation = robot.get_observation() + assert observation["J1.pos"] == 0.0 + assert observation["J6.pos"] == 5.0 + + robot.disconnect() + assert arm.mode == 0 + assert ("disconnect",) in arm.calls + + +def test_manual_mode_config_rejects_cartesian_control(tmp_path): + with pytest.raises(ValueError, match="control_space='joint'"): + UFRobotConfig( + id="test_manual_robot", + calibration_dir=tmp_path, + robot_dof=6, + control_space="cartesian", + manual_mode=True, + ) + + +def test_manual_action_filters_non_action_observation_fields(): + observation = { + "J1.pos": 1.0, + "J1.vel": 2.0, + "gripper.pos": 0.5, + "camera": np.zeros((2, 2, 3), dtype=np.uint8), + } + action_features = {"J1.pos": float, "gripper.pos": float} + + assert _manual_action_from_observation(observation, action_features) == { + "J1.pos": 1.0, + "gripper.pos": 0.5, + } + + +def test_manual_record_config_has_no_teleop(monkeypatch): + config_path = Path("config/manual_mode/xarm7_manual_record_config.yaml").resolve() + monkeypatch.setattr( + sys, + "argv", + ["uf-lerobot-record", "--config_path", str(config_path)], + ) + + config = get_cfg() + + assert config.robot.manual_mode is True + assert config.robot.robot_dof == 7 + assert config.teleop is None + assert config.dataset.fps == 30 + + +def test_manual_record_loop_writes_actual_state_as_action(tmp_path): + class FakeRobot: + name = "fake_manual_robot" + robot_type = name + action_features = {"J1.pos": float, "J2.pos": float} + observation_features = action_features + + def __init__(self): + self.observation_count = 0 + self.sent_actions = [] + + def get_observation(self): + self.observation_count += 1 + value = float(self.observation_count) + return {"J1.pos": value, "J2.pos": value + 1} + + def send_action(self, action): + self.sent_actions.append(action.copy()) + return action + + robot = FakeRobot() + action_pipeline, robot_pipeline, observation_pipeline = record_module.make_default_processors() + features = record_module.combine_feature_dicts( + record_module.aggregate_pipeline_dataset_features( + pipeline=action_pipeline, + initial_features=record_module.create_initial_features(action=robot.action_features), + use_videos=False, + ), + record_module.aggregate_pipeline_dataset_features( + pipeline=observation_pipeline, + initial_features=record_module.create_initial_features( + observation=robot.observation_features + ), + use_videos=False, + ), + ) + dataset = LeRobotDataset.create( + "test/manual-record", + fps=30, + features=features, + root=tmp_path / "dataset", + robot_type=robot.robot_type, + use_videos=False, + ) + + record_module.record_loop( + robot=robot, + events={"exit_early": False}, + fps=30, + teleop_action_processor=action_pipeline, + robot_action_processor=robot_pipeline, + robot_observation_processor=observation_pipeline, + dataset=dataset, + control_time_s=0.001, + single_task="test task", + manual_mode=True, + ) + + assert dataset.episode_buffer["size"] == 1 + assert robot.sent_actions == [{"J1.pos": 2.0, "J2.pos": 3.0}] + assert dataset.episode_buffer["action"][0].tolist() == [2.0, 3.0] + + dataset.save_episode() + dataset.finalize()