rfm_so101_collection.py

from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
from lerobot.cameras.opencv.camera_opencv import OpenCVCamera
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.utils import hw_to_dataset_features
from lerobot.robots.so101_follower import SO101Follower, SO101FollowerConfig
from lerobot.teleoperators.so101_leader.config_so101_leader import SO101LeaderConfig
from lerobot.teleoperators.so101_leader.so101_leader import SO101Leader
from lerobot.utils.control_utils import init_keyboard_listener
from lerobot.utils.utils import log_say
from lerobot.utils.visualization_utils import init_rerun
from lerobot.scripts.lerobot_record import record_loop
from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig
from lerobot.cameras.configs import ColorMode, Cv2Rotation
from lerobot.processor import make_default_processors

import json
import argparse
import time
import sys, select

curr_time = time.time()
TOP_CAM_INDEX = 4
WRIST_CAM_INDEX = 6


def main(args):
    dataset_name = args.dataset_name
    dataset_name = dataset_name + str(curr_time)


    fps = args.fps
    episode_time_sec = args.episode_time_sec
    reset_time_sec = args.reset_time_sec

    # Create the robot and teleoperator configurations
    # camera_config = {
    #     "front": RealSenseCameraConfig(
    #         serial_number_or_name="239222302862",
    #         fps=fps,
    #         width=640,
    #         height=480,
    #         color_mode=ColorMode.RGB,
    #         use_depth=False,
    #         rotation=Cv2Rotation.NO_ROTATION
    #     )
    # }
    # camera_config = {
    camera_config = {
        "top": OpenCVCameraConfig(index_or_path=TOP_CAM_INDEX, width=640, height=480, fps=30),
        "wrist": OpenCVCameraConfig(index_or_path=WRIST_CAM_INDEX, width=640, height=480, fps=30),
    }


    # camera = OpenCVCamera(camera_config)
    # camera.connect()

    robot_config = SO101FollowerConfig(
        port="/dev/ttyACM2", id="so101_follower_arm", cameras=camera_config
    )

    # Initialize the robot and teleoperator
    teleop_config = SO101LeaderConfig(port="/dev/ttyACM1", id="so101_leader_arm")
    robot = SO101Follower(robot_config)
    teleop = SO101Leader(teleop_config)
    teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()

    # Configure the dataset features
    action_features = hw_to_dataset_features(robot.action_features, "action")
    obs_features = hw_to_dataset_features(robot.observation_features, "observation")
    dataset_features = {**action_features, **obs_features}





    num_episodes = args.num_episodes

    dataset = LeRobotDataset.create(
        repo_id="HenryZhang/" + dataset_name,
        fps=fps,
        features=dataset_features,
        robot_type=robot.name,
        use_videos=True,
        image_writer_threads=2,
    )

    _, events = init_keyboard_listener()
    init_rerun(session_name="recording")

    with open(args.ann_json_path, "r") as f:
        annotations = json.load(f)

    log_say("Connecting robot and teleop...")
    robot.connect()
    teleop.connect()

    try:
        for TASK_DESCRIPTION in annotations:
            log_say(f"Resetting the arm, Starting task: {TASK_DESCRIPTION}. Setting the environment accordingly. Press ENTER to continue.")

            # Small teleop manual warm-up loop
            while True:
                teleop_action = teleop.get_action()
                robot.send_action(teleop_action)
                time.sleep(0.01)

                # Check ENTER press
                if sys.stdin in select.select([sys.stdin], [], [], 0)[0]:
                    input()
                    break

            episode_idx = 0
            while episode_idx < num_episodes and not events["stop_recording"]:
                log_say(f"Recording episode {episode_idx + 1} of {num_episodes}")

                record_loop(
                    robot=robot,
                    events=events,
                    fps=fps,
                    teleop=teleop,
                    dataset=dataset,
                    control_time_s=episode_time_sec,
                    single_task=TASK_DESCRIPTION,
                    display_data=True,
                    teleop_action_processor=teleop_action_processor,
                    robot_action_processor=robot_action_processor,
                    robot_observation_processor=robot_observation_processor,
                )

                # Reset phase
                if not events["stop_recording"] and (episode_idx < num_episodes - 1 or events["rerecord_episode"]):
                    log_say("Reset the environment")
                    record_loop(
                        robot=robot,
                        events=events,
                        fps=fps,
                        teleop=teleop,
                        control_time_s=reset_time_sec,
                        single_task=TASK_DESCRIPTION,
                        display_data=True,
                        teleop_action_processor=teleop_action_processor,
                        robot_action_processor=robot_action_processor,
                        robot_observation_processor=robot_observation_processor,
                    )

                if events["rerecord_episode"]:
                    log_say("Re-recording episode")
                    events["rerecord_episode"] = False
                    events["exit_early"] = False
                    dataset.clear_episode_buffer()
                    continue

                dataset.save_episode()
                episode_idx += 1

    except KeyboardInterrupt:
        log_say("❗ Ctrl-C detected — stopping early and disconnecting safely.")

    finally:
        log_say("Stopping recording and disconnecting devices...")
        try:
            robot.disconnect()
        except:
            pass
        try:
            teleop.disconnect()
        except:
            pass

        # # Only push if normal exit (no Ctrl-C interruption before completion)
        # if not isinstance(sys.exc_info()[1], KeyboardInterrupt):
        #     log_say("Pushing dataset to hub...")
        #     try:
        #         dataset.push_to_hub()
        #         log_say("✅ Push complete!")
        #     except Exception as e:
        #         log_say(f"⚠️ Failed to push dataset: {e}")
        # else:
        #     log_say("Dataset NOT pushed because recording was interrupted.")

        # log_say("✅ episode_time_secShutdown finished.")
        base_path = "/home/koch/.cache/huggingface/lerobot/HenryZhang/" + dataset_name

        # import subprocess
        # import time

        # cmd = f"""
        # hf upload HenryZhang/{dataset_name} . . \
        # --repo-type=dataset \
        # --delete="*"
        # """

        # print("Running command:")
        # print(cmd)

        # time.sleep(3)  # wait for filesystem to settle

        # subprocess.run(
        #     cmd,
        #     cwd=base_path,   # run inside dataset folder
        #     shell=True,      # EXACT shell behavior
        #     check=True,      # crash if upload fails
        # )

        # print("Upload complete.", "dataset_name:", dataset_name)
        print("please manually upload the dataset located at:", base_path)
        print("command is:", f"hf upload HenryZhang/{dataset_name} . --repo-type=dataset --delete='*'")



        # from huggingface_hub import HfApi
        # api = HfApi()
        # api.create_repo(
        #     repo_id=f"HenryZhang/{dataset_name}",
        #     repo_type="dataset",
        #     exist_ok=True,
        # )

        # api.upload_large_folder(
        #     folder_path=str(base_path),
        #     repo_id=f"HenryZhang/{dataset_name}",
        #     repo_type="dataset",
        # )



if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Record data for SO101 robot")
    parser.add_argument(
        "--ann_json_path", type=str, default="vla_tasks.json", help="Task annotation JSON"
    )
    parser.add_argument(
        "--num_episodes", type=int, default=40, help="Number of episodes to record"
    )
    parser.add_argument(
        "--dataset_name", type=str, default="RFM_fail", required=True, help="Name of the dataset to create"
    )
    parser.add_argument(
        "--fps", type=int, default=15, help="Frames per second for recording"
    )
    parser.add_argument(
        "--episode_time_sec", type=int, default=20, help="Run in headless mode without visualization"
    )
    parser.add_argument(
        "--reset_time_sec", type=int, default=7, help="Reset time between episodes"
    )
    args = parser.parse_args()

    main(args)