RoboTwin2.0-RGB-D

RoboTwin 2.0 RGBD Dataset

中文文档

An RGBD robotic manipulation dataset regenerated from RoboTwin 2.0, augmented with depth camera data for training imitation learning policies such as PI-0.

Uses the Large_D435 camera configuration (640×480 resolution), matching the raw shape in configs/data/robotwin.yaml.

Dataset Description

Both versions cover all 50 manipulation tasks in RoboTwin 2.0, with depth: true and Large_D435 cameras (640×480).

Cameras

Each camera provides both RGB and depth (uint16, millimeters) data.

Dataset Structure

LeRobot Format

Each frame contains:

Depth values are stored as raw uint16 millimeters — no compression, no scaling. 0 indicates invalid depth.

Raw HDF5 Format

Raw collected data is stored as HDF5 files. Each episode contains:

• joint_action/ — left/right arm joint positions + gripper state
• observation/<camera>/rgb — RGB images (JPEG encoded)
• observation/<camera>/depth — depth maps (float32, millimeters)
• instructions/ — natural language task instructions

Reproduction Guide

The following describes how to regenerate this dataset. Data collection and processing scripts are located in robotwin_depth_patch/, which can be rsync'd into the RoboTwin repo root:

rsync -av robotwin_depth_patch/ /path/to/RoboTwin/

1. Install Environment

Dev machine image: pytorch:2.8.0-cu12.9.1-py3.10-ubuntu22.04-0.1

Follow the official installation guide, then install additional dependencies:

pip install sapien==3.0.0b1 open3d

Installation Tips:

• In scripts/_install.sh, comment out pytorch3d and curobo installation commands (install curobo manually instead);
• For Download Assets, download the archive from HF Web first, then upload to the VS Code directory;
• If you encounter a blinker uninstall error, run pip install blinker --upgrade --ignore-installed blinker before installing open3d;
• Install ffmpeg: apt update && apt install -y ffmpeg && ffmpeg -version

2. Create Configuration Files

bash create_config.sh

Generates Large_D435 depth configs from demo_clean.yml and demo_randomized.yml:

3. Collect Data

Single Task Collection (for testing)

# Format: bash collect_data.sh <task_name> <task_config> <gpu_id>

# clean version (50 episodes)
bash collect_data.sh beat_block_hammer demo_clean_largeD435_depth 0

# randomized version (500 episodes)
bash collect_data.sh beat_block_hammer demo_randomized_largeD435_depth 0

Batch Task Collection

cd /home/hongxiao_yu/RoboTwin

# Format: bash collect_data_all.sh <task_config> <gpu_csv> [tasks_per_gpu] [extra_args...]

# ===== clean version (50 episodes/task) =====

bash collect_data_all.sh demo_clean_largeD435_depth 0                           # 1 GPU serial (most stable, recommended for initial validation)
bash collect_data_all.sh demo_clean_largeD435_depth 0,1,2,3,4,5,6,7 1          # 8 GPUs × 1 serial
bash collect_data_all.sh demo_clean_largeD435_depth 0,1,2,3,4,5,6,7 2          # 8 GPUs × 2 parallel

# ===== randomized version (500 episodes/task) =====

bash collect_data_all.sh demo_randomized_largeD435_depth 0
bash collect_data_all.sh demo_randomized_largeD435_depth 0,1,2,3,4,5,6,7 1
bash collect_data_all.sh demo_randomized_largeD435_depth 0,1,2,3,4,5,6,7 2

# ===== Common options =====

bash collect_data_all.sh demo_clean_largeD435_depth 0 1 adjust_bottle  # Single task (others skipped, supports resumption)
bash collect_data_all.sh demo_clean_largeD435_depth 0 --force           # Force restart (stops existing instance first)

Status: bash scripts/status_collect.sh ｜ Stop: bash scripts/stop_collect.sh

# Vulkan errors
apt install -y libglvnd0 libegl1 libvulkan1 mesa-vulkan-drivers
mkdir -p /usr/share/glvnd/egl_vendor.d

# open3d errors
apt install -y libgl1-mesa-glx libglfw3

# Bus error
pip install open3d==0.17.0

# Additional dependencies
pip install toppra mplib

4. Generate LeRobot Format Data

Follow the openpi data guide. The original PI-0 pipeline does not support depth data, so the following files must be added first (all included in robotwin_depth_patch/):

Run the conversion:

cd policy/pi0

# clean version (50 episodes)
bash process_data_pi0_depth.sh beat_block_hammer demo_clean_largeD435_depth 50
bash generate_with_depth.sh ./training_data/demo_clean_largeD435_depth/ demo_clean_largeD435_depth_repo

# randomized version (500 episodes)
bash process_data_pi0_depth.sh beat_block_hammer demo_randomized_largeD435_depth 500
bash generate_with_depth.sh ./training_data/demo_randomized_largeD435_depth/ demo_randomized_largeD435_depth_repo

Output Directory Control

process_data_pi0_depth.sh accepts an optional 4th argument for the output base directory (defaults to /mnt/data_volume2/data2/RoboTwin2.0/processed_data), forwarded to process_data_with_depth.py --out_dir:

# Single task with custom output dir -> <out_dir>/<task>-<config>-<num>/episode_*/episode_*.hdf5
bash process_data_pi0_depth.sh beat_block_hammer demo_clean_largeD435_depth 50 /path/to/your/processed_data

Batch Processing (All Tasks)

The patch also ships a tmux-based batch entry at policy/pi0/process_all_pi0_depth.sh (included in robotwin_depth_patch/):

cd policy/pi0

# Format: bash process_all_pi0_depth.sh <task_config> <expert_data_num> [parallel] [extra_tasks...]
# - 1 tmux session + N worker windows (each worker serially runs its assigned tasks)
# - Round-robin task assignment; done tasks auto-skipped (resume-friendly)
# - Per-task logs: <OUT_DIR>/.logs/<task>.log  |  failed: <OUT_DIR>/.failed  |  partial: <OUT_DIR>/.partial
# - Broken episodes are recorded in <task-dir>/.broken and subtracted from the expected count

# Process all 50 tasks with 8 parallel workers
bash process_all_pi0_depth.sh demo_clean_largeD435_depth 50 8

# Process only specific tasks
bash process_all_pi0_depth.sh demo_clean_largeD435_depth 50 8 adjust_bottle beat_block_hammer

# Attach to view real-time progress (status window + per-worker windows)
tmux attach -t <session_name>

# Status (without attaching)
bash script/process_status.sh [session_name]

# Stop
bash script/process_stop.sh [session_name]

Note: process_all_pi0_depth.sh hardcodes OUT_DIR=/mnt/data_volume2/data2/RoboTwin2.0/processed_data_3. Edit this variable inside the script to change the output location. It passes OUT_DIR as the 4th argument to process_data_pi0_depth.sh, which forwards it as --out_dir to process_data_with_depth.py.

5. Other Related Scripts in the Repo

The scripts below are referenced by the pipeline above. The first two (data collection) are now included in robotwin_depth_patch/; the rest are setup/eval scripts that already live in the rsync'd repo and are documented here for completeness.

Setup & Configuration

Data Collection (Python Entries)

Both scripts are now included in robotwin_depth_patch/script/. Use collect_broken_episode.py to repair episodes that failed with UnStableError or Render Error after a batch run, instead of re-running the whole task. Requires the seed and trajectory pkl for that index to already exist.

Policy Training & Evaluation (PI-0)

Distributed Evaluation (Client / Server)

For multi-GPU eval, run the policy as a TCP server and let eval clients connect to it:

# 1) Start model server (GPU node) — loads checkpoint once, serves infer requests over TCP
python script/policy_model_server.py --config policy/pi0/deploy_policy.yml --port <port> \
    --overrides --task_name <task> --task_config <config> \
    --train_config_name <train_cfg> --model_name <model> --ckpt_setting <model> \
    --policy_name pi0

# 2) Run eval client (runs the Sapien sim, connects to the server for actions)
python script/eval_policy_client.py --config policy/pi0/deploy_policy.yml --port <port> \
    --overrides --task_name <task> --task_config <config> \
    --train_config_name <train_cfg> --model_name <model> --ckpt_setting <model> \
    --policy_name pi0

Alternatively, openpi ships a websocket policy server:

uv run policy/pi0/scripts/serve_policy.py --config <train_config_name> --dir <checkpoint_dir>

Non-depth Variants (Reference)

The non-depth counterparts of the depth pipeline above (kept for reference; do not use them for this RGBD dataset):

6. Failure Handling & Re-collection / Re-conversion

Episodes can be skipped at three stages. Below is how to detect and repair each.

6.1 Data Collection (collect_data.py / collect_data_all.sh)

Task-level skip / fail — tracked by marker files under logs/:

# Check status
bash scripts/status_collect.sh          # prints done=N fail=M total=50 + current running tasks

# Retry a single failed task
rm -f logs/<task>.fail
bash collect_data_all.sh demo_clean_largeD435_depth 0 1 <task>

# One-click retry ALL failed tasks
for t in $(awk '{print $1}' logs/_failed.log); do rm -f logs/$t.fail; done
bash collect_data_all.sh demo_clean_largeD435_depth 0

# Force restart the whole batch (stops existing instance first)
bash collect_data_all.sh demo_clean_largeD435_depth 0 --force

Episode-level skip during collection — collect_data.py resumes from the first missing data/episodeN.hdf5 (by index), so re-running the same task auto-continues. Two failure modes:

# Re-collect a specific broken episode (keeps good episodes untouched)
python script/collect_broken_episode.py <task> <task_config> 481
# Re-collect multiple at once
python script/collect_broken_episode.py <task> <task_config> 481 100 200

Prerequisite: seed.txt and _traj_data/episodeN.pkl for that index must already exist (they do, unless you deleted them).

Find missing episodes (compare seed count vs hdf5 count):

cd data/<task>/<task_config>
echo "seeds:   $(wc -w < seed.txt)"
echo "hdf5s:   $(ls data/episode*.hdf5 2>/dev/null | wc -l)"
# List missing indices
python -c "import os; n=int(open('seed.txt').read().split().__len__()); \
print([i for i in range(n) if not os.path.isfile(f'data/episode{i}.hdf5')])"

6.2 Data Processing (process_data_with_depth.py / process_all_pi0_depth.sh)

Task-level skip / partial / fail — tracked under <OUT_DIR>/:

# Check batch status (without attaching tmux)
bash script/process_status.sh [session_name]

# Re-run the whole batch — done tasks auto-skip, partial/failed get re-processed
bash process_all_pi0_depth.sh demo_clean_largeD435_depth 50 8

# Re-run only specific tasks
bash process_all_pi0_depth.sh demo_clean_largeD435_depth 50 8 adjust_bottle beat_block_hammer

# Re-run a single task directly (bypass tmux)
bash process_data_pi0_depth.sh <task> <config> <num> <out_dir>

Episode-level broken during processing — process_data_with_depth.py skips and records broken indices:

.broken indices are subtracted from the expected count by process_all_pi0_depth.sh, so a task with 50 target and 2 broken is considered done at 48 episodes. To re-process a broken episode after fixing the source:

cd <OUT_DIR>/<task>-<config>-<num>
# 1. Remove the broken marker (so it counts toward expected again)
rm -f .broken
# 2. Remove the broken episode's output dir (so it's not skipped on resume)
rm -rf episode_481/
# 3. Re-run processing for this task
bash ../../process_data_pi0_depth.sh <task> <config> <num> <OUT_DIR>

6.3 LeRobot Conversion (generate_with_depth.sh)

If the conversion is interrupted, just re-run the same command — the converter re-scans the source episode_*/ dirs and the LeRobot dataset writer appends/overwrites per-episode. No special cleanup needed:

bash generate_with_depth.sh ./training_data/<config>/ <repo_id>

To start fresh (re-convert everything from scratch), delete the cached dataset first:

rm -rf ~/.cache/huggingface/lerobot/<repo_id>/
bash generate_with_depth.sh ./training_data/<config>/ <repo_id>