HeRO

Official implementation of HeRO: Hierarchical 3D Semantic Representation for Pose-aware Object Manipulation.

Chongyang Xu, Shen Cheng, Haipeng Li, Haoqiang Fan, Ziliang Feng, Shuaicheng Liu

HeRO learns pose-aware 3D semantic representations for dual-arm object manipulation in RoboTwin. This repository provides the policy implementation, RoboTwin task adapters used by the GroundedSAM and tracking pipeline, and scripts for data collection, preprocessing, training, and evaluation.

News

• 2026-06-08: Code, RoboTwin task adapters, preprocessing, training, and evaluation scripts are prepared for release.
• 2026-02-21: The paper is available on arXiv.
• 2026-01: HeRO is accepted to ICRA 2026.

Highlights

• Standalone policy repository that does not require modifying the external RoboTwin codebase.
• Adapted RoboTwin task layer for GroundedSAM-based segmentation and tracking.
• End-to-end reproduction scripts covering data collection, feature preprocessing, training, and evaluation.
• Configurable local paths for GroundingDINO, BERT, SAM, DINOv2, and CleanDIFT weights.
• Evaluation outputs are written to a local results/ directory by default.

Repository Layout

.
├── collect_data.sh              # Collect RoboTwin demonstrations
├── process_data.sh              # Build zarr datasets and PCA models
├── train.sh                     # Train the HeRO policy
├── eval.sh                      # Evaluate a trained checkpoint
├── deploy_policy.yml            # Evaluation-time policy configuration
├── hero/                        # HeRO policy, model, training, and runtime code
│   ├── HeRO.py                   # Policy import entry for POLICY_NAME=HeRO
│   ├── config/                  # Hydra training configs
│   ├── dataset/                 # zarr dataset loader
│   ├── env_runner/              # Policy rollout runner
│   ├── model/                   # Neural network modules
│   ├── policy/                  # HeRO policy implementation
│   └── runtime/                 # Deployment, tracking, and feature utilities
├── pose_aware_tasks/            # RoboTwin task adapters and configs
├── pretrained/                  # Weight download scripts
├── scripts/                     # Python drivers used by root shell scripts
└── thirdparty/                  # Lightweight local wrappers

The repository loads pose_aware_tasks before the external ROBOTWIN_ROOT through PYTHONPATH/sys.path. These files adapt existing RoboTwin tasks for the GroundedSAM and tracking pipeline while keeping the external RoboTwin checkout unchanged.

Requirements

Prepare a working RoboTwin environment first, then clone this repository independently. The scripts assume:

• ROBOTWIN_ROOT points to the external RoboTwin root.
• RoboTwin assets are available under ${ROBOTWIN_ROOT}/assets.
• Grounded-Segment-Anything exposes Detect_and_Seg.py.
• CleanDIFT is importable as a Python package.
• SAM checkpoint is available through SAM_CHECKPOINT or the default RoboTwin third-party path.
• DINOv2 is available through torch.hub, or a local repo is specified by DINO_REPO_DIR.

HeRO preprocessing requires raw demonstrations with:

• observation/head_camera/depth
• pointcloud

The included pose_aware_tasks/task_config/demo_clean.yml enables these fields.

Pretrained Weights

Model weights are intentionally not tracked by git. Download GroundingDINO, BERT, and CleanDIFT weights with:

bash pretrained/download_weights.sh

For slower Hugging Face access:

HF_ENDPOINT=https://hf-mirror.com bash pretrained/download_weights.sh

The default HeRO path only needs CompVis/cleandift's cleandift_sd21_full.safetensors. You may also download it manually:

mkdir -p pretrained/CleanDIFT
huggingface-cli download CompVis/cleandift cleandift_sd21_full.safetensors \
  --local-dir pretrained/CleanDIFT

If Hugging Face requires authentication:

huggingface-cli login

or:

HF_TOKEN=<your_huggingface_token> huggingface-cli download CompVis/cleandift cleandift_sd21_full.safetensors \
  --local-dir pretrained/CleanDIFT

Then set the weight paths:

export GROUNDINGDINO_CONFIG=/path/to/HeRO/pretrained/GroundingDINO/GroundingDINO_SwinB.cfg.py
export GROUNDINGDINO_CHECKPOINT=/path/to/HeRO/pretrained/GroundingDINO/groundingdino_swinb_cogcoor.pth
export BERT_BASE_UNCASED_PATH=/path/to/HeRO/pretrained/GroundingDINO/bert-base-uncased
export CLEANDIFT_CHECKPOINT=/path/to/HeRO/pretrained/CleanDIFT/cleandift_sd21_full.safetensors

HeRO can construct the CleanDIFT VAE and UNet from cleandift_sd21_full.safetensors; Stable Diffusion 2.1 is only needed when running the upstream CleanDIFT DiffusionPipeline or base-model branch.

Minimal Reproduction

Run from the HeRO repository root inside the RoboTwin Python environment:

export ROBOTWIN_ROOT=/path/to/RoboTwin-icra
export GROUNDINGDINO_CONFIG=/path/to/HeRO/pretrained/GroundingDINO/GroundingDINO_SwinB.cfg.py
export GROUNDINGDINO_CHECKPOINT=/path/to/HeRO/pretrained/GroundingDINO/groundingdino_swinb_cogcoor.pth
export BERT_BASE_UNCASED_PATH=/path/to/HeRO/pretrained/GroundingDINO/bert-base-uncased
export CLEANDIFT_CHECKPOINT=/path/to/HeRO/pretrained/CleanDIFT/cleandift_sd21_full.safetensors

bash collect_data.sh place_phone_stand_g demo_clean 0
bash process_data.sh place_phone_stand_g demo_clean 100 5 1024 0
bash train.sh place_phone_stand_g demo_clean 100 0 0
bash eval.sh place_phone_stand_g demo_clean demo_clean 100 300 0 "0"

The command sequence performs:

1. Data collection with the adapted RoboTwin task layer.
2. HeRO feature extraction, PCA fitting, and zarr dataset creation.
3. Policy training for the selected task and dataset split.
4. RoboTwin evaluation with the trained checkpoint.

Expected generated paths:

pose_aware_tasks/data/place_phone_stand_g/demo_clean/
data/zarr_data/place_phone_stand_g_demo_clean_100_1024_5_HeRO.zarr
PCA_Model/place_phone_stand_g_demo_clean_100_1024_5_HeRO/
checkpoints/place_phone_stand_g_demo_clean_100_1024_5_HeRO_0/place_phone_stand_g-HeRO-pose_query/300.ckpt
results/place_phone_stand_g/HeRO/demo_clean/place_phone_stand_g-HeRO-pose_query/seed_0/300/_result.txt

Useful Environment Variables

• ROBOTWIN_ROOT: external RoboTwin root.
• ROBOTWIN_TASK_ROOT: adapted RoboTwin task root, default pose_aware_tasks.
• RAW_DATA_ROOT: raw demonstration root, default pose_aware_tasks/data.
• OUTPUT_ROOT: zarr output root, default data/zarr_data.
• PCA_ROOT: PCA output root, default PCA_Model.
• DATA_ROOT / ZARR_PATH: training data override.
• CKPT_PATH: evaluation checkpoint override.
• RESULTS_ROOT: evaluation result root, default results.
• TEST_NUM: number of evaluation episodes, default 100.
• NUM_EPOCHS: training epochs, default 300.
• CHECKPOINT_EVERY: checkpoint interval, default 100.
• PYTHON_BIN: Python executable, default python3.

Acknowledgements

This repository builds on RoboTwin, G3Flow, and CleanDIFT. We thank the authors and maintainers of these projects for their open-source contributions.