Diagnosing Compositional Generalization in Sequential Robot Tasks

Overview

Robots do not need demonstrations of every instruction combination. A structured quarter of the task space generalizes to the rest, and when one step constrains another, coverage must span those dependencies.

1/4 → 78.2%

held-out success using
a quarter of the task space

0.4% → 54.7%

success on pretraining-held-out tasks
after 1 demo per instruction

3.4×

held-out success from spreading
a fixed budget across dependencies

Abstract

Sequential robot manipulation requires policies to execute novel combinations of familiar instruction components. However, collecting demonstrations for all possible instruction tuples is combinatorially expensive, while sparsely covered datasets often fail under out-of-distribution recombination. This paper studies compositional generalization through the lens of instruction-space coverage. We decompose the generalization gap into three sources: marginal instruction shift, instruction-compositional shift, and context–action shift. Our results show that exhaustive tuple enumeration is unnecessary: a structured subset, as small as one quarter of the full task space, can recover strong out-of-distribution performance when it covers action-relevant dependencies.

Generalization Framework

Where does the generalization gap come from? Our analysis separates three sources of distribution shift.

1. Marginal Instruction Shift

Whether individual instruction values have similar coverage and frequency in training and test data.

2. Instruction-Compositional Shift

Whether familiar subtasks are recombined in unfamiliar ways at test time.

3. Context–Action Shift

How observation–action distributions differ even after conditioning on the full instruction sequence.

Experiments

Q1 Sparse Task Coverage

Does the robot need to be trained on every task?

Pick–Place–Press has 4 objects × 4 receptacles × 4 buttons = 64 tasks. We train on B of them and test on the held-out rest.

Pick–Place–Press · 64 tasks · open standalone ↗

The orthogonal 16

Every pair of factor values appears exactly once.

object receptacle button

Success rate vs. B

Takeaway: A structured quarter of the task space reaches 78.2% held-out success, matching full training (78.6%).

Every task, one at a time

Open full screen ↗

Each cube is one of the 64 tasks. Pick a training set and hover to see per-task success.

Q2 Instruction Finetuning

What is missing when sparse coverage fails?

We pretrain on B tasks, then finetune with one demonstration per instruction, and test on the tasks held out from pretraining.

Pick–Place–Press · 64 tasks · open standalone ↗

Pretrained only

+ 1 demo / instruction

OOD success before and after finetuning

Takeaway: Sparse pretraining already learns the skills; one demonstration per instruction steers them to new combinations (0.4% → 54.7%).

Q3 Dependent Instructions

Which combinations must training cover when one step constrains another?

The robot is told where the second object goes (c₂); the first object may go anywhere else (c₁ ≠ c₂).

Dependent two-stage Pick-and-Place · 48 tasks · open standalone ↗

The task: “Put the cube away, then put the milk on the plate.”

Three frames of a robot arm. Step 1: pick up the cube — blue arrows point to the purple cup, red cup and box as valid destinations, a red arrow to the plate is crossed out. Step 2: pick up the milk. Step 3: place the milk to the plate. A red arrow loops back to step 1, captioned 'Cannot put the cube into plate!', showing that the stage-2 target constrains stage 1.

Same 16 training tasks, spread over more (o₁, c₂) combinations

Every condition trains on B = 16 tasks; only the number of covered (o₁, c₂) combinations changes.

Combinations covered by the same 16 tasks

c₂ — reserved for the second object

trained held out c₁ = c₂, forbidden

16 / 16

combinations trained

legal first-placement cases covered

B = 16

training tasks — same in every condition

held-out tasks

Task success

± SE

Takeaway: At a fixed budget of 16 tasks, covering more dependent combinations raises held-out success from 9.2% to 31.4%.

Further Findings

Stage-wise modular policy vs. full-instruction policy

Stage-wise modular π(a | z, l_σ(z)) A ground-truth stage indicator passes only the active instruction.

Full instruction π(a | z, l) The policy sees every instruction at every stage.

(a) Independent instructions: the active instruction is enough

Pick–Place–Press · OOD success rate (%)

(b) Dependent instructions: the hidden constraint is lost

Two-stage Pick-and-Place · wrong first container (%), lower is better

In (b), the first placement must avoid the later container c₂, which is not part of the active instruction. Cov16 covers all 16 dependent pairs; Full trains on all 48 tasks. Error bars follow Figure 9 of the paper.

Takeaway: Stage-wise modularity helps when instructions are independent, but cannot learn a dependency on an instruction it does not see.

Code & Data

The training code, simulator, and demonstration datasets for all three task families are public. Setup, data preparation, and train/evaluate commands are in the code repository README.

Code Diagnosing_Compositional_Generalization_Robot RoboPolicy learning library plus the paper's data preparation, training, and evaluation tools. Simulator robot_envs MuJoCo environments for Pick-and-Place, Pick–Place–Press, and the dependent two-stage task. Dataset cg_l4_1s Pick-and-Place (PP) demonstrations: 4 objects × 4 receptacles. Dataset cg_l5 Pick–Place–Press (PPP) demonstrations: 64 object–receptacle–button instructions. Dataset cg_l4_2s Dependent two-stage Pick-and-Place (2S-PP) demonstrations; both two-stage variants use this dataset.

Datasets are hosted on Hugging Face. Pretrained checkpoint downloads are not yet available.

Citation

Accepted to the Conference on Robot Learning (CoRL) 2026 as a spotlight presentation. Preprint: arXiv:2607.29687.

@inproceedings{wang2026diagnosing,
  title={Diagnosing Compositional Generalization in Sequential Robot Tasks},
  author={Yixiao Wang and Cheng-En Wu and Lingfeng Sun and Pengcheng Wang and
          Xiang Ji and Boyuan Liang and Guojian Zhan and Masayoshi Tomizuka},
  booktitle={Conference on Robot Learning (CoRL)},
  year={2026},
  note={Spotlight presentation},
  url={https://arxiv.org/abs/2607.29687}
}

01 Overview