Paper Rob Rdd

Method Overview

RDD addresses planner-visuomotor dataset misalignment in hierarchical VLA frameworks. Instead of relying on human annotations or heuristics (like UVD) for sub-task decomposition, RDD:

1. Embeds all frames of a demonstration using a pre-trained visual encoder (LIV/CLIP/R3M)
2. Retrieves similar sub-task segments from a prior database via approximate nearest neighbor (ANN) search
3. Decomposes the demonstration via dynamic programming that finds the optimal partition maximizing retrieval similarity, subject to length constraints

Key insight: Formulate sub-task decomposition as an optimal partitioning problem where the score of each candidate segment is its similarity to the nearest prior sub-task — achieving near-oracle performance (only 0.2% success rate gap vs expert annotations) with linear time complexity.

Paper-Code Mapping

Setup

Dependencies

• Python 3.9
• PyTorch (with CUDA)
• annoy>=1.17.0 (Spotify ANN library)
• fastapi>=0.115.0, uvicorn>=0.34.0
• sqlmodel>=0.0.24, confection>=0.1.0, typer>=0.15.0
• loguru, pyzmq, elara, scipy, scikit-learn, opencv-python, imageio
• External: UVD (github.com/zcczhang/UVD), LIV (github.com/penn-pal-lab/LIV)

Installation

# Create environment
conda create -n rdd python==3.9
conda activate rdd

# Clone repo
git clone https://github.com/tasl-lab/Retrieval-Demonstration-Decomposer.git
cd Retrieval-Demonstration-Decomposer

# Install Python deps
pip install -r scripts/setup/requirements.txt

# Clone external model libraries into 3rdparty/
bash scripts/setup/setup_rdd_env.sh
# This clones UVD (e7d657a) and LIV (a12991f) into 3rdparty/

Architecture

Input: Long-horizon demonstration (T frames, V views)
  │
  ├─ 1. EMBED: frames -> visual encoder (LIV/CLIP/R3M)
  │      Output: embeds (T, D) per view
  │      Multi-view: concat across views -> (T, D*V)
  │
  ├─ 2. PRIOR DATABASE (built offline from training sub-tasks):
  │      For each annotated sub-task -> extract feature
  │        default mode: concat(start_frame, end_frame) -> (2D*V,)
  │        ood mode: end_frame only -> (D*V,)
  │      Store in Annoy ANN index
  │
  ├─ 3. SCORE each candidate segment [s, e]:
  │      a. Extract sub-task feature -> query ANN -> retrieval_dist
  │      b. Length penalty: alpha * |1 - (e-s) / nn.duration|
  │      c. UVD penalty: beta * |s - uvd_start| / (e-s)
  │      d. score = -(retrieval_dist + penalties) * segment_len
  │
  ├─ 4. DYNAMIC PROGRAMMING: max_sum_partition()
  │      dp[i] = max over j of (dp[j] + score(u[j:i]))
  │      subject to min_len <= i-j <= max_len
  │
  └─ 5. OUTPUT: optimal partition -> list of sub-task segments
         + matched prior sub-tasks (nearest neighbors)

Key Hyperparameters

Usage Scenarios

1. Build Vector Database (offline, one-time)

python build_vec_database.py <gpu_id> <encoder> <sample_rate> <dataset_path> \
    --name-suffix <name> --views front_rgb --embed-mode ood

2. Run RDD Server

uvicorn rdd_server:app --port 8001 --workers 8

Endpoints:

• GET /config — returns current config
• POST /decompose — accepts frame paths, returns segments + nearest neighbors

3. Evaluate Decomposition Quality

python eval_rdd.py <dataset_path> <results_save_path> \
    --worker-num 4 --rdd-port 8001

Computes per-segment IoU against ground truth annotations.

Code Integration Guide

Minimal Imports

from rdd.algorithms import max_sum_partition, rdd_score
from rdd.datasets.rlbench import RLBenchAnnoySearcher
from rdd.embed import uvd_embed, subtask_embeds_to_feature

Decompose a New Demonstration

import numpy as np

# 1. Embed frames
embeds = uvd_embed(
    frame_paths,                    # list of image file paths
    preprocessor="liv",             # or "clip", "r3m"
    device="cuda",
)  # -> (T, D) np.ndarray

# 2. Load prior sub-task database + ANN index
searcher = RLBenchAnnoySearcher(
    searcher_path="data/vec_databases/franka/train/index.ann",
    vec_database_path="data/vec_databases/franka/train",
    include_views=["front_rgb"],
    n_trees=10,
    distance_measure="angular",
)

# 3. Run optimal partitioning
score, segments = max_sum_partition(
    u=list(range(len(embeds))),
    score_func=rdd_score,
    min_len=2,
    max_len=100,
    searcher=searcher,
    embeds=embeds,
    mode="ood",
    alpha=0.0,          # forced to 0 in ood mode
    beta=0.1,
)
# segments: list of lists, e.g. [[0,1,...,29], [30,31,...,74], ...]

Data Format

Supported Encoders

Repo Structure

Key Results

• Only 0.2% success rate gap compared to expert-annotated (oracle) decompositions
• Outperforms UVD heuristic decomposition on both simulation (RLBench) and real-world tasks
• Training-free — no model training required, only pre-trained visual encoders + ANN indexing
• Linear time complexity via the optimal partitioning DP algorithm

Tips & Gotchas

• RDD is training-free — no custom weights to train. The "prior" is built from your annotated sub-task dataset via build_vec_database.py
• The mode parameter matters: ood (end-frame only, alpha=0) works better for out-of-distribution tasks; default (start+end, with length penalty) is for in-distribution
• UVD and LIV must be cloned into 3rdparty/ via setup_rdd_env.sh — they are not pip-installable
• The server architecture (FastAPI) is designed for integration with a higher-level planner that sends decomposition requests
• beta controls how much RDD defers to UVD's heuristic boundaries — set higher (0.25) for tasks where UVD works well, lower (0.1) for novel tasks
• Multi-view support: views are concatenated in embedding space, increasing feature dimension proportionally
• info.txt annotation format: each line is the ending frame index of a sub-task (not the starting index)
• Annoy indexes are not updatable — rebuild with build_vec_database.py when adding new sub-tasks