From 47c531889b009bdcd356a0cf1ccc346e649de7ac Mon Sep 17 00:00:00 2001 From: ShashwatPatil Date: Tue, 1 Jul 2025 00:37:03 +0530 Subject: [PATCH] added readme, dataset recording guide and ACT policy training for so101 arm --- examples/lerobot-dataset-record/Readme.md | 137 ++++++++++++++++++ examples/lerobot-dataset-record/dataflow.yml | 12 +- node-hub/dora-dataset-record/README.md | 103 ++++++++++--- .../dora_dataset_record/main.py | 1 + node-hub/dora-dataset-record/pyproject.toml | 3 +- 5 files changed, 230 insertions(+), 26 deletions(-) create mode 100644 examples/lerobot-dataset-record/Readme.md diff --git a/examples/lerobot-dataset-record/Readme.md b/examples/lerobot-dataset-record/Readme.md new file mode 100644 index 00000000..38d347d4 --- /dev/null +++ b/examples/lerobot-dataset-record/Readme.md @@ -0,0 +1,137 @@ +# Dataset Recording with SO101 Robot Arm + +Guide for recording imitation learning datasets with the SO101 arm using Dora and train policies with LeRobot. + +## Overview + +To begin recording, edit the dataflow.yml file. You must update camera settings, robot ports, and dataset configuration based on your setup. + +### Required Changes + +Update the camera device id, you can also add or remove cameras based on your setup: + +```yaml +# Laptop camera +CAPTURE_PATH: "0" # Laptop camera is defaults to 0 +# External camera +CAPTURE_PATH: "1" # Change this to match your external camera device +``` +### SO101 Arms + +Identify and set the correct USB ports for both the leader and follower SO101 arms +```yaml +PORT: "/dev/ttyACM0" # change this +``` +### Dataset Recorder Configuration + +>Edit the following fields: + +```yaml +REPO_ID: "your_username/so101_dataset" # HuggingFace dataset path +SINGLE_TASK: "Pick up the red cube and place it in the blue box" # Your task description +CAMERA_NAMES: "laptop, front" # Name your camera sources depending on your setup +CAMERA_LAPTOP_RESOLUTION: "480,640,3" +CAMERA_FRONT_RESOLUTION: "480,640,3" +``` + +You can adjust the following parameters as needed: + +```yaml +TOTAL_EPISODES: "2" # Number of episodes +#you may want to try with 2-3 episodes to test, then atleast 50 episodes for training is recommended +EPISODE_DURATION_S: "60" # Duration of each episode (in seconds) - depends on complexity of task +RESET_DURATION_S: "15" # Time to reset the environment between episodes +FPS: "30" # Should match camera fps +PUSH_TO_HUB: "false" # Set to "true" to auto-upload dataset to HuggingFace +ROOT_PATH: "full path where you want to save the dataset" +# if not defined then it will be stored at ~/.cache/huggingface/lerobot/repo_id +``` + +Once everything is updated in `dataflow.yml`, you are ready to record your dataset. + +## Start Recording +Build and Run + +```bash +dora build dataflow.yml +dora run dataflow.yml +``` + +## Recording Process + +In the rerun window you can see the the info regarding Start of episodes, start of reset phase and saving of episodes. + +#### During Recording + +1. **Episode Active**: + - Use the **leader robot** to demonstrate/perform the task + - Move smoothly and naturally + - Complete the full task within the time limit + +2. **Reset Phase**: + - move the leader arm to initial position + - Reset objects to starting positions + - Prepare workspace for next demonstration + +3. **Repeat** until all episodes are complete + +**Recording Tips** + +- **Practice First**: Move the leader robot around to get comfortable before recording +- **Smooth Movements**: Avoid jerky or sudden movements +- **Complete Tasks**: Try to finish the task within each episode +- **Consistent Setup**: Keep object positions and lighting consistent +- **Monitor Rerun**: Watch the visualization to ensure cameras and robot data are flowing + +## After Recording + +Your dataset will be saved locally. Check the recording was successful: +>It will be stored in ~/.cache/huggingface/lerobot/repo_id + + +# Training and Testing Policies using the recorded dataset + +After successfully recording your dataset, we will be training imitation learning policies and deploying them on your SO101 robot arm. + +## Training Your Policy + +#### Install LeRobot Training Dependencies + +Easiest way to train your policy is to use lerobots training scripts +```bash +# Install training requirements and lerobot repo +git clone https://github.com/huggingface/lerobot.git +pip install lerobot[training] +pip install tensorboard wandb # For monitoring (Optional) +``` + +### Choose Your Policy + +**ACT (Recommended for SO101)** +- Good for manipulation tasks +- Handles multi-modal data well +- Faster training (for me it took 7hrs 🥺 to train on 50 episodes for pick and place, 3050 laptop gpu) + +**Diffusion Policy** +- Better for complex behaviors +- More robust to distribution shift +- Longer training + + +### Start Training + +```bash +cd lerobot + +python lerobot/scripts/train.py \ + --dataset.repo_id=${HF_USER}/your_repo_id \ # provide full path of your dataset + --policy.type=act \ + --output_dir=outputs/train/act_so101_test \ + --job_name=act_so101_test \ + --policy.device=cuda \ + --wandb.enable=true \ + --policy.repo_id=${HF_USER}/my_policy +``` + +You can monitor your training progress on wandb +> For more details regarding training check [lerobot](https://huggingface.co/docs/lerobot/en/il_robots#train-a-policy) guide on Imitation learning for SO101 and testing your policy \ No newline at end of file diff --git a/examples/lerobot-dataset-record/dataflow.yml b/examples/lerobot-dataset-record/dataflow.yml index a8476190..c734c2fc 100644 --- a/examples/lerobot-dataset-record/dataflow.yml +++ b/examples/lerobot-dataset-record/dataflow.yml @@ -20,7 +20,7 @@ nodes: outputs: - image env: - CAPTURE_PATH: "2" + CAPTURE_PATH: "1" ENCODING: "rgb8" IMAGE_WIDTH: "640" IMAGE_HEIGHT: "480" @@ -48,7 +48,7 @@ nodes: PORT: "/dev/ttyACM0" IDS: "1,2,3,4,5,6" - - id: lerobot_recorder + - id: dataset_recorder build: pip install -e ../../node-hub/dora-dataset-record path: dora-dataset-record inputs: @@ -59,12 +59,12 @@ nodes: outputs: - text env: - REPO_ID: "HF_username/name_of_dataset" + REPO_ID: "HF_username/Name_your_dataset" SINGLE_TASK: "Pick up the cube and place it in the box" - ROBOT_TYPE: "your robot type" # e.g., "so101_follower", "ur5e" etc + ROBOT_TYPE: "so101_follower" # e.g., "koch", "franka", "ur5e" etc FPS: "30" - TOTAL_EPISODES: "5" + TOTAL_EPISODES: "50" EPISODE_DURATION_S: "60" RESET_DURATION_S: "15" @@ -85,4 +85,4 @@ nodes: inputs: image_laptop: laptop_cam/image image_front: front_cam/image - text: lerobot_recorder/text \ No newline at end of file + text: dataset_recorder/text \ No newline at end of file diff --git a/node-hub/dora-dataset-record/README.md b/node-hub/dora-dataset-record/README.md index 9aa61889..d48fa763 100644 --- a/node-hub/dora-dataset-record/README.md +++ b/node-hub/dora-dataset-record/README.md @@ -1,40 +1,107 @@ # dora-dataset-record -## Getting started +Node for recording robot datasets in LeRobot format. You can captures synchronized camera feeds and robot poses to create high-quality datasets for imitation learning and robot training. -- Install it with uv: +- **Robot pose recording** - Capture both state and action data +- **Multi-camera support** - Record from multiple cameras simultaneously +- **LeRobot dataset format (v2.1)** - Direct integration with HuggingFace LeRobot datasets +- **Episode management** - Automatic episode segmentation with reset phases + +## Quick Start + +### 1. Installation ```bash -uv venv -p 3.11 --seed +# Source your venv +cd dora/node-hub/dora-dataset-record uv pip install -e . ``` -## Contribution Guide +### 2. Usage Guide -- Format with [ruff](https://docs.astral.sh/ruff/): +Create a dataflow file, see `examples/lerobot-dataset-record/dataflow.yml`: -```bash -uv pip install ruff -uv run ruff check . --fix +```yaml +nodes: + # Dataset recorder + - id: dataset_recorder + build: pip install -e ../../node-hub/dora-dataset-record + path: dora-dataset-record + inputs: + laptop: laptop_cam/image + front: front_cam/image + robot_state: robot_follower/pose + robot_action: leader_interface/pose + outputs: + - text + env: + # Required settings + REPO_ID: "your_username/your_dataset_name" + SINGLE_TASK: "Pick up the cube and place it in the box" + ROBOT_TYPE: "your_robot_type" + + # Recording settings + FPS: "30" + TOTAL_EPISODES: "50" + EPISODE_DURATION_S: "60" + RESET_DURATION_S: "15" + + # Camera configuration + CAMERA_NAMES: "laptop,front" + CAMERA_LAPTOP_RESOLUTION: "480,640,3" + CAMERA_FRONT_RESOLUTION: "480,640,3" + + # Robot configuration + ROBOT_JOINTS: "joint1,joint2,joint3,joint4,joint5,gripper" + + # Optional settings + USE_VIDEOS: "true" + PUSH_TO_HUB: "false" + PRIVATE: "false" + TAGS: "robotics,manipulation,imitation_learning" + + # Visualization with rerun + - id: plot + build: pip install dora-rerun + path: dora-rerun + inputs: + text: dataset_recorder/text ``` -- Lint with ruff: +### 3. Start Recording the dataset ```bash -uv run ruff check . +dora build dataflow.yml +dora run dataflow.yml ``` -- Test with [pytest](https://github.com/pytest-dev/pytest) +The node will send instructions on dora-rerun, about episode starting, reset time, Saving episodes etc. -```bash -uv pip install pytest -uv run pytest . # Test -``` +## Configuration + +### Required Environment Variables + +| Variable | Description | Example | +|----------|-------------|---------| +| `REPO_ID` | HuggingFace dataset repo | `"username/dataset_name"` | +| `SINGLE_TASK` | Task description | `"Pick and place objects"` | +| `CAMERA_NAMES` | Comma-separated camera names | `"laptop,front,top"` | +| `CAMERA_*_RESOLUTION` | Resolution for each camera | `"480,640,3"` | +| `ROBOT_JOINTS` | Comma-separated joint names | `"joint1,joint2,gripper"` | -## YAML Specification +### Optional Settings -## Examples +| Variable | Default | Description | +|----------|---------|-------------| +| `FPS` | `30` | Recording frame rate (match camera fps) | +| `TOTAL_EPISODES` | `10` | Number of episodes to record | +| `EPISODE_DURATION_S` | `60` | Episode length in seconds | +| `RESET_DURATION_S` | `15` | Break between episodes to reset the environment | +| `USE_VIDEOS` | `true` | Encode as MP4 videos, else saves images | +| `PUSH_TO_HUB` | `false` | Upload to HuggingFace Hub | +| `PRIVATE` | `false` | Make dataset private | +| `ROOT_PATH` | `~/.cache/huggingface/lerobot/your_repo_id` | Local storage path where you want to save the dataset | ## License -dora-dataset-record's code are released under the MIT License +This project is released under the MIT License. \ No newline at end of file diff --git a/node-hub/dora-dataset-record/dora_dataset_record/main.py b/node-hub/dora-dataset-record/dora_dataset_record/main.py index 5617b8db..e03f39ce 100644 --- a/node-hub/dora-dataset-record/dora_dataset_record/main.py +++ b/node-hub/dora-dataset-record/dora_dataset_record/main.py @@ -311,6 +311,7 @@ class DoraLeRobotRecorder: """Output message.""" # Put message in queue to send self.message_queue.put(message) + print(message) def get_pending_messages(self): """Get all pending messages from the queue.""" diff --git a/node-hub/dora-dataset-record/pyproject.toml b/node-hub/dora-dataset-record/pyproject.toml index 8cf43844..5854596a 100644 --- a/node-hub/dora-dataset-record/pyproject.toml +++ b/node-hub/dora-dataset-record/pyproject.toml @@ -27,5 +27,4 @@ ignore = [ ] [tool.uv.sources] -lerobot = {git = "https://github.com/huggingface/lerobot.git" } - +lerobot = {git = "https://github.com/huggingface/lerobot.git", rev = "main"} \ No newline at end of file