RDD is an retrieval-based visual demonstration decomposer that automatically identifies sub-tasks visually similar to a set of existing expert-labeled sub-tasks.

Mingxuan Yan¹, Yuping Wang^1,2, Zechun Liu³, Jiachen Li¹

¹ University of California, Riverside
² University of Michigan
³ Meta AI

• Sub-task Discovery with Prior: Different to non-prior heuristic sub-task discovery algorithms such as UVD, RDD identifies sub-tasks that are visually similar to ones in a given expert labeled sub-task dataset. This is specially useful when generating additional sub-tasks for fine-tuning or data augmentation, which encourages the policy to reuse learned skills from the original dataset.
• Planner-visuomotor Alignment: (Youtube Video) In hierarchical VLAs, the planner, often a powerful VLM, performs task planning and reasoning to break down complex tasks into simpler sub-tasks with step-by-step language instructions. Conditioned on the generated sub-task instructions, a learning-based visuomotor policy, trained on datasets with short-horizon sub-tasks, performs precise manipulation to complete the sub-tasks one by one, thereby completing long-horizon tasks. RDD automatically decomposes demonstrations into sub-tasks by aligning the visual features of the decomposed sub-task intervals with those from the training data of the low-level visuomotor policies.

RDD formulates demonstration decomposition as an optimal partitioning problem, using retrieval with approximate nearest neighbor search (ANNS) and dynamic programming to efficiently find the optimal decomposition strategy.

Set up python environment:

conda create -n rdd python==3.9 -y && conda activate rdd
./scripts/setup/setup_rdd_env.sh

This default installation only supports encoders LIV, ClIP, ResNet. To install full support for other encoders ( VIP , R3M , VC-1) please follow setup_rdd_env.sh.

See kitchen_demo.md

See franka_demo.md

See agi_cerebra_demo.md

See racer_demo.md

See reproduction.md for which experiment reproduces each table, the compute and disk budget, and per-experiment seed counts.

RDD allows you to define your own sub-task prior.

In RDD the sub-task is represented by the feature of of ending frame / starting frame. You can define your own sub-task feature by modifying sub-task feature generation class subtask_embeds_to_feature in rdd/embed.py

class subtask_embeds_to_feature(object):
	@staticmethod
	def feature_dim(embed_dim: int, mode: str) -> int:
		<your feature dim>
	def __call__(self, embeds: Union[np.ndarray, torch.Tensor], mode: str) -> np.ndarray:
		"""
		input: embeds (L, N)
		output: feature (2N,) or (N,) if mode=='ood'
		"""
		<your implementation>

You can also define a completely custimized sub-task prior score by modifying rdd_score in rdd/algorithms.py

def rdd_score(
	subarray: List[int], 
	searcher: AnnoySearcher, 
	embeds: np.ndarray, 
	...
	):
	<your implementation>

• Release the core algorithm and demo scripts on AgiBotWorld & RoboCerebra. (ETA: by the end of Oct.2025).
• Release use case scripts on hiearchical VLA (RACER).

RDD builds on a number of open-source projects and datasets, and we thank their authors for releasing their code and data.

Decomposition. UVD serves both as a baseline and, via rdd/uvd_wrapper.py, as a component of RDD's scoring function.

Visual encoders. LIV, R3M, VIP, VC-1, CLIP, and DINOv2. Retrieval uses Annoy.

Hierarchical VLA use case (Example #4). Built on RACER and its data-generation pipeline, which in turn build on RVT-2 and PerAct. The planner is LLaVA-NeXT, finetuned and served through Open-LLaVA-NeXT. Evaluation runs in RLBench on PyRep/CoppeliaSim.

Datasets. AgiBotWorld-Alpha and RoboCerebra (Example #3), and RLBench demonstrations (Example #4).

Vendored utility. scripts/setup/switch_cuda.sh is the MIT-licensed switch-cuda utility, Copyright (c) 2018 Patrick Hohenecker, included unmodified with its license header.

If you find this work useful, please consider citing our paper:

@inproceedings{yan2025rdd,
  title={RDD: Retrieval-Based Demonstration Decomposer for Planner Alignment in Long-Horizon Tasks},
  author={Yan, Mingxuan and Wang, Yuping and Liu, Zechun and Li, Jiachen},
  booktitle={Proceedings of the 39th Annual Conference on Neural Information Processing Systems (NeurIPS)},
  year={2025},
}