feat: add SO101 recording and ACT workflow with progress tracking

2026-09-10 16:46:58 +08:00
parent 2e33355a9c
commit cbf51facff
16 changed files with 604 additions and 7 deletions
@@ -0,0 +1,26 @@
+#!/usr/bin/env bash
+# Shared task configuration; source from a launcher.
+set -eo pipefail
+TASK_SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+source "$TASK_SCRIPT_DIR/../env/activate_lerobot.sh"
+set -u
+TASK_TEXT='Pick up the block and place it in the white box.'
+DATASET_ID="${DATASET_ID:-local/block_white_box_v1}"
+DATASET_ROOT="${DATASET_ROOT:-$LEROBOT_PROJECT_DIR/data/block_white_box_v1}"
+TASK_CAMERAS='{"wrist": {"type": "opencv", "index_or_path": "/dev/v4l/by-id/usb-ZSKJ-260319-K_Integrated_Webcam_HD-video-index0", "width": 640, "height": 480, "fps": 30, "fourcc": "MJPG"}, "front": {"type": "opencv", "index_or_path": "/dev/v4l/by-id/usb-Sonix_Technology_Co.__Ltd._USB2.0_CAM1_USB2.0_CAM1-video-index0", "width": 640, "height": 480, "fps": 25, "fourcc": "MJPG"}}'
+ROBOT_ARGS=(--robot.type=so101_follower "--robot.port=$FOLLOWER_PORT" --robot.id=zihao_follower_arm "--robot.cameras=$TASK_CAMERAS")
+run_task() {
+  if [[ "${DRY_RUN:-0}" == 1 ]]; then
+    printf '%q ' "$@"
+    printf '\n'
+    return
+  fi
+  exec "$@"
+}
+check_task_devices() {
+  [[ "${DRY_RUN:-0}" == 1 ]] && return 0
+  local device
+  for device in "$FOLLOWER_PORT" "/dev/v4l/by-id/usb-ZSKJ-260319-K_Integrated_Webcam_HD-video-index0" "/dev/v4l/by-id/usb-Sonix_Technology_Co.__Ltd._USB2.0_CAM1_USB2.0_CAM1-video-index0"; do
+    [[ -r "$device" && -w "$device" ]] || { echo "设备未连接或缺少读写权限：$device" >&2; exit 1; }
+  done
+}
@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/block_white_box_common.sh"
+if [[ "${RESUME:-false}" == true ]]; then
+  [[ -f "$DATASET_ROOT/meta/info.json" ]] || { echo "找不到可续录的数据：$DATASET_ROOT" >&2; exit 1; }
+else
+  [[ ! -e "$DATASET_ROOT" ]] || { echo "目录已存在；续录请设置 RESUME=true，或指定新的 DATASET_ROOT / DATASET_ID。" >&2; exit 1; }
+fi
+check_task_devices
+if [[ "${DRY_RUN:-0}" != 1 ]]; then
+  [[ -r "$LEADER_PORT" && -w "$LEADER_PORT" ]] || { echo "主臂未连接或缺少权限：$LEADER_PORT" >&2; exit 1; }
+fi
+printf '\n========== 积木放入白盒：准备采集 ==========\n'
+echo "本次 ${NUM_EPISODES:-5} 条；每条最多 ${EPISODE_SECONDS:-30} 秒；复位最多 ${RESET_SECONDS:-30} 秒。"
+echo "绿色横幅 = 正在录制；黄色横幅 = 复位 / 重录。"
+echo "→ 完成当前阶段    ← 重录当前条    Esc 结束采集"
+echo "请在此终端操作方向键，并留意阶段横幅。"
+if [[ "${DRY_RUN:-0}" != 1 ]]; then
+  read -r -p "摆好积木、盒子和主从臂后，按 Enter 连接设备并开始录制：" _ready
+fi
+run_task python "$TASK_SCRIPT_DIR/record_with_prompts.py" "${ROBOT_ARGS[@]}" \
+  --teleop.type=so101_leader "--teleop.port=$LEADER_PORT" --teleop.id=zihao_leader_arm \
+  "--dataset.repo_id=$DATASET_ID" "--dataset.root=$DATASET_ROOT" \
+  "--dataset.single_task=$TASK_TEXT" --dataset.no_stamp=true \
+  --dataset.fps=25 "--dataset.num_episodes=${NUM_EPISODES:-5}" \
+  "--dataset.episode_time_s=${EPISODE_SECONDS:-30}" "--dataset.reset_time_s=${RESET_SECONDS:-30}" \
+  --dataset.push_to_hub=false --dataset.rgb_encoder.vcodec=h264 \
+  --display_data=true "--resume=${RESUME:-false}" "$@"
@@ -0,0 +1,67 @@
+#!/usr/bin/env python3
+"""Chinese phase banners for LeRobot recording; no changes to robot control."""
+import sys
+
+
+def show_phase(message: str) -> None:
+    if message.startswith("Recording episode "):
+        number = int(message.rsplit(" ", 1)[1]) + 1
+        title, color = f"● 正在录制｜数据集第 {number} 条", "32"
+        lines = [
+            "现在操作主臂：抓起积木 → 放入白盒 → 松开 → 退开",
+            "完成：按 → 结束本条，进入复位（最后一条将保存退出）",
+            "失败：按 ← 标记重录，然后按提示复位",
+            "停止：按 Esc；当前片段可能保存，之后检查完整性",
+        ]
+    elif message == "Reset the environment":
+        title, color = "● 复位阶段｜当前不录入示范", "33"
+        lines = [
+            "把积木取出放回起点，并通过主臂让从臂回到起始姿态",
+            "注意：从臂仍跟随主臂！",
+            "准备好：按 → 继续；超时也会自动继续",
+            "上一条失败：按 ← 标记重录；按 Esc 结束采集",
+        ]
+    elif message == "Re-record episode":
+        title, color = "● 已丢弃本条｜即将重录", "33"
+        lines = ["等待绿色“正在录制”提示后，再开始抓取。"]
+    elif message == "Stop recording":
+        title, color = "● 正在结束采集｜请等待保存与清理", "36"
+        lines = ["不要关闭终端；完成后会返回命令行。"]
+    elif message == "Exiting":
+        title, color = "● 采集程序已退出", "36"
+        lines = ["请检查实际保存条数及视频；提前结束的片段可能不完整。"]
+    else:
+        return
+    use_color = sys.stdout.isatty()
+    start = f"\033[1;{color}m" if use_color else ""
+    end = "\033[0m" if use_color else ""
+    print("\n" + start + "━" * 60, flush=True)
+    print(title, flush=True)
+    print("━" * 60, flush=True)
+    for line in lines:
+        print("  " + line, flush=True)
+    print("━" * 60 + end + "\n", flush=True)
+    if use_color:
+        print("\a", end="", flush=True)
+
+
+def main() -> None:
+    if sys.argv[1:] == ["--preview"]:
+        for message in ["Recording episode 0", "Reset the environment", "Re-record episode", "Stop recording", "Exiting"]:
+            show_phase(message)
+        return
+    from lerobot.scripts import lerobot_record
+
+    original_log_say = lerobot_record.log_say
+
+    def log_with_banner(message, *args, **kwargs):
+        result = original_log_say(message, *args, **kwargs)
+        show_phase(message)
+        return result
+
+    lerobot_record.log_say = log_with_banner
+    lerobot_record.main()
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,13 @@
+#!/usr/bin/env bash
+source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/block_white_box_common.sh"
+if [[ -z "${MODEL_DIR:-}" ]]; then
+  echo "请设置 MODEL_DIR 为检查点的 pretrained_model 目录。" >&2; exit 1
+fi
+if [[ "${DRY_RUN:-0}" != 1 && ! -f "$MODEL_DIR/config.json" ]]; then
+  echo "模型配置不存在：$MODEL_DIR/config.json" >&2; exit 1
+fi
+check_task_devices
+echo "自主执行 ${DURATION:-30} 秒；按 Ctrl+C 停止。"
+run_task lerobot-rollout --strategy.type=base "${ROBOT_ARGS[@]}" \
+  "--policy.path=$MODEL_DIR" --policy.device=cuda "--task=$TASK_TEXT" \
+  --fps=25 "--duration=${DURATION:-30}" --display_data=true "$@"
@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+# Launch SO101 teleoperation with this machine's calibrated arms.
+set -eo pipefail
+
+script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+source "$script_dir/../env/activate_lerobot.sh"
+set -u
+
+check_device() {
+  if [[ ! -e "$1" ]]; then
+    echo "设备未连接：$1" >&2
+    exit 1
+  fi
+  if [[ ! -r "$1" || ! -w "$1" ]]; then
+    echo "设备缺少读写权限：$1，请检查串口 dialout / 摄像头 video 权限。" >&2
+    exit 1
+  fi
+}
+check_device "$FOLLOWER_PORT"
+check_device "$LEADER_PORT"
+
+echo "启动遥操作；按 Ctrl+C 停止。"
+exec lerobot-teleoperate \
+  --robot.type=so101_follower \
+  --robot.port="$FOLLOWER_PORT" \
+  --robot.id=zihao_follower_arm \
+  --robot.cameras='{}' \
+  --teleop.type=so101_leader \
+  --teleop.port="$LEADER_PORT" \
+  --teleop.id=zihao_leader_arm \
+  --display_data=false \
+  "$@"
@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+# Launch SO101 teleoperation with this machine's calibrated arms.
+set -eo pipefail
+
+script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+source "$script_dir/../env/activate_lerobot.sh"
+set -u
+
+check_device() {
+  if [[ ! -e "$1" ]]; then
+    echo "设备未连接：$1" >&2
+    exit 1
+  fi
+  if [[ ! -r "$1" || ! -w "$1" ]]; then
+    echo "设备缺少读写权限：$1，请检查串口 dialout / 摄像头 video 权限。" >&2
+    exit 1
+  fi
+}
+check_device "$FOLLOWER_PORT"
+check_device "$LEADER_PORT"
+check_device "/dev/v4l/by-id/usb-ZSKJ-260319-K_Integrated_Webcam_HD-video-index0"
+
+echo "启动遥操作；按 Ctrl+C 停止。"
+exec lerobot-teleoperate \
+  --robot.type=so101_follower \
+  --robot.port="$FOLLOWER_PORT" \
+  --robot.id=zihao_follower_arm \
+  --robot.cameras='{"wrist": {"type": "opencv", "index_or_path": "/dev/v4l/by-id/usb-ZSKJ-260319-K_Integrated_Webcam_HD-video-index0", "width": 640, "height": 480, "fps": 30, "fourcc": "MJPG"}}' \
+  --teleop.type=so101_leader \
+  --teleop.port="$LEADER_PORT" \
+  --teleop.id=zihao_leader_arm \
+  --display_data=true \
+  "$@"
@@ -0,0 +1,34 @@
+#!/usr/bin/env bash
+# Launch SO101 teleoperation with this machine's calibrated arms.
+set -eo pipefail
+
+script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+source "$script_dir/../env/activate_lerobot.sh"
+set -u
+
+check_device() {
+  if [[ ! -e "$1" ]]; then
+    echo "设备未连接：$1" >&2
+    exit 1
+  fi
+  if [[ ! -r "$1" || ! -w "$1" ]]; then
+    echo "设备缺少读写权限：$1，请检查串口 dialout / 摄像头 video 权限。" >&2
+    exit 1
+  fi
+}
+check_device "$FOLLOWER_PORT"
+check_device "$LEADER_PORT"
+check_device "/dev/v4l/by-id/usb-ZSKJ-260319-K_Integrated_Webcam_HD-video-index0"
+check_device "/dev/v4l/by-id/usb-Sonix_Technology_Co.__Ltd._USB2.0_CAM1_USB2.0_CAM1-video-index0"
+
+echo "启动遥操作；按 Ctrl+C 停止。"
+exec lerobot-teleoperate \
+  --robot.type=so101_follower \
+  --robot.port="$FOLLOWER_PORT" \
+  --robot.id=zihao_follower_arm \
+  --robot.cameras='{"wrist": {"type": "opencv", "index_or_path": "/dev/v4l/by-id/usb-ZSKJ-260319-K_Integrated_Webcam_HD-video-index0", "width": 640, "height": 480, "fps": 30, "fourcc": "MJPG"}, "front": {"type": "opencv", "index_or_path": "/dev/v4l/by-id/usb-Sonix_Technology_Co.__Ltd._USB2.0_CAM1_USB2.0_CAM1-video-index0", "width": 640, "height": 480, "fps": 25, "fourcc": "MJPG"}}' \
+  --teleop.type=so101_leader \
+  --teleop.port="$LEADER_PORT" \
+  --teleop.id=zihao_leader_arm \
+  --display_data=true \
+  "$@"
@@ -0,0 +1,120 @@
+#!/usr/bin/env python3
+"""Move one calibrated SO-101 follower joint a few degrees and return."""
+
+import argparse
+import os
+import time
+from pathlib import Path
+
+from lerobot.robots.so_follower import SO101Follower, SO101FollowerConfig
+
+
+JOINTS = (
+    "shoulder_pan",
+    "shoulder_lift",
+    "elbow_flex",
+    "wrist_flex",
+    "wrist_roll",
+    "gripper",
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--port",
+        default=os.environ.get("FOLLOWER_PORT"),
+        help="Follower serial port; defaults to FOLLOWER_PORT from env/activate_lerobot.sh.",
+    )
+    parser.add_argument("--robot-id", default="my_follower", help="ID used when calibrating the follower.")
+    parser.add_argument("--joint", choices=JOINTS, default="wrist_roll")
+    parser.add_argument("--delta", type=float, default=3.0, help="Signed movement in degrees (gripper: percent).")
+    parser.add_argument("--duration", type=float, default=1.0, help="Seconds for each outbound/return motion.")
+    return parser.parse_args()
+
+
+def interpolate(robot: SO101Follower, joint: str, start: float, delta: float, duration: float) -> None:
+    steps = max(2, round(duration * 20))
+    for step in range(1, steps + 1):
+        target = start + delta * step / steps
+        robot.bus.write("Goal_Position", joint, target)
+        time.sleep(duration / steps)
+
+
+def main() -> None:
+    args = parse_args()
+    if not args.port:
+        raise SystemExit("Missing port. Run: source env/activate_lerobot.sh")
+    port = Path(args.port)
+    if not port.exists():
+        raise SystemExit(f"Follower port is not connected: {port}")
+    if not os.access(port, os.R_OK | os.W_OK):
+        raise SystemExit(
+            f"No read/write permission for {port}. Add the user to dialout, then log in again or run newgrp dialout."
+        )
+    if not 0 < abs(args.delta) <= 5:
+        raise SystemExit("For this first-motion test, --delta must be non-zero and at most 5.")
+    if not 0.5 <= args.duration <= 5:
+        raise SystemExit("--duration must be between 0.5 and 5 seconds.")
+
+    config = SO101FollowerConfig(port=str(port), id=args.robot_id)
+    robot = SO101Follower(config)
+    if not robot.calibration_fpath.is_file():
+        raise SystemExit(
+            f"Calibration file not found: {robot.calibration_fpath}\n"
+            "Calibrate first with lerobot-calibrate, using the same --robot.id."
+        )
+
+    print(f"Calibration: {robot.calibration_fpath}")
+    print(f"Test: {args.joint} {args.delta:+g}, then return to the measured start position.")
+    print("Place the arm in a supported pose, clear the workspace, and keep power cutoff within reach.")
+    if input("Type MOVE and press Enter to continue: ").strip() != "MOVE":
+        raise SystemExit("Cancelled; no movement command was sent.")
+
+    selected_torque_enabled = False
+    try:
+        # Use the motor bus directly: SOFollower.connect() configures and enables
+        # every motor, which is inappropriate for a first single-joint test.
+        robot.bus.connect()
+        robot.bus.disable_torque()
+        if not robot.bus.is_calibrated:
+            raise RuntimeError("Motor calibration does not match the local calibration file.")
+
+        start = float(robot.bus.read("Present_Position", args.joint))
+        print(f"Start {args.joint}.pos: {start:.2f}")
+        # Set a non-moving goal before enabling torque, avoiding a stale goal
+        # register from a previous session.
+        robot.bus.write("Goal_Position", args.joint, start)
+        robot.bus.enable_torque(args.joint)
+        selected_torque_enabled = True
+
+        interpolate(robot, args.joint, start, args.delta, args.duration)
+        time.sleep(0.3)
+        outbound = float(robot.bus.read("Present_Position", args.joint))
+        actual_delta = outbound - start
+        print(f"Outbound {args.joint}.pos: {outbound:.2f} (actual {actual_delta:+.2f})")
+        if actual_delta * args.delta <= 0 or abs(actual_delta) < abs(args.delta) * 0.5:
+            raise RuntimeError(
+                f"Joint feedback did not follow the command: requested {args.delta:+.2f}, observed {actual_delta:+.2f}."
+            )
+
+        interpolate(robot, args.joint, start + args.delta, -args.delta, args.duration)
+        time.sleep(0.3)
+        returned = float(robot.bus.read("Present_Position", args.joint))
+        return_error = returned - start
+        print(f"Returned {args.joint}.pos: {returned:.2f} (error {return_error:+.2f})")
+        if abs(return_error) > 1.0:
+            raise RuntimeError(f"Joint did not return within 1 degree of start; error is {return_error:+.2f}.")
+        print("Motion test completed and verified by position feedback.")
+    finally:
+        if robot.bus.is_connected:
+            try:
+                if selected_torque_enabled:
+                    print(f"Disabling torque on {args.joint}.")
+                    robot.bus.disable_torque(args.joint, num_retry=5)
+            finally:
+                robot.bus.disconnect(disable_torque=False)
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,14 @@
+#!/usr/bin/env bash
+source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/block_white_box_common.sh"
+if [[ "${DRY_RUN:-0}" != 1 && ! -f "$DATASET_ROOT/meta/info.json" ]]; then
+  echo "请先录制并检查数据：$DATASET_ROOT" >&2; exit 1
+fi
+if [[ "${DRY_RUN:-0}" != 1 ]]; then
+  python -c 'import torch; import sys; sys.exit(0 if torch.cuda.is_available() else "CUDA 不可用，请先修复 NVIDIA 驱动；已停止训练以避免回退到 CPU。")'
+fi
+TRAIN_OUTPUT="${TRAIN_OUTPUT:-$LEROBOT_PROJECT_DIR/outputs/train/act_block_white_box_$(date +%Y%m%d_%H%M%S)}"
+run_task lerobot-train --policy.type=act --policy.device=cuda --policy.push_to_hub=false \
+  "--dataset.repo_id=$DATASET_ID" "--dataset.root=$DATASET_ROOT" \
+  "--output_dir=$TRAIN_OUTPUT" --job_name=act_block_white_box \
+  "--batch_size=${BATCH_SIZE:-8}" "--steps=${STEPS:-20000}" --num_workers=4 \
+  --save_freq=2000 --env_eval_freq=0 --seed=42 --wandb.enable=false "$@"