Add manual mode data collection script

.gitignore

@@ -88,3 +88,4 @@ models/
 *.trt
 *.xvcd
 ufactory_usage/
+.history/

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.
 

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. 推理
 
 指定模型进行推理
 

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

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 遥操作

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:

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")

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):
@@ -189,6 +213,39 @@ class AsyncEpisodeSaver:
             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
 def record_loop(
     robot: Robot,
@@ -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:
             },
         )
 
+    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,8 +574,9 @@ 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()
+                    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)
@@ -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():

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()