Diagnosing Compositional Generalization in Sequential Robot Tasks

Yixiao Wang1, Cheng-En Wu1, Lingfeng Sun1, Pengcheng Wang1, Xiang Ji2, Boyuan Liang1, Guojian Zhan2, Masayoshi Tomizuka1

1University of California, Berkeley 2Tsinghua University

Paper (PDF) Code Dataset Supplementary

TL;DR

Full tuple coverage is unnecessary — a structured sparse subset recovers strong OOD performance, and one demo per instruction closes most of what remains. What decides which sparse subset works is our third result: when one step constrains another, coverage must be spread across those dependencies. At a fixed budget, spreading the same 16 tasks over more dependency cases more than triples held-out success — so the design rule is not how much data to collect, but along which axis to spread it.

1/4 → 78.2%
a quarter of the task space
recovers full OOD success
0.4% → 54.7%
OOD success after
1-demo-per-instruction finetune
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.

Where the Generalization Gap Comes From

We bound the compositional generalization gap by three terms, each with a different interpretation and controllability.

1. Marginal Instruction Shift

Whether individual subtask instructions are covered by training. This is a data issue resolvable only by expanding training support.

2. Instruction-Compositional Shift

Whether familiar subtasks are recombined in unfamiliar ways at test time. Pairwise coverage in the training set controls this term.

3. Context–Action Shift

The residual shift in context–action behavior when the policy is overly sensitive to irrelevant instruction correlations.

Experiments

Q1: Does the robot need to be trained on every task?

In Pick-Place-Press, one instruction means: pick one of 4 objects, place it in one of 4 receptacles, press one of 4 buttons. That is 4 × 4 × 4 = 64 possible instructions.

We train on B of them and test on all 64. The other 64 − B are held out: the same skills, in a combination never shown. Held-out success climbs steeply with B, then plateaus at B = 16 — a quarter of the set is enough.

And which 16 matters. The orthogonal 16 on the left covers every pair of ingredients exactly once, and reaches that plateau by itself.

Left: how the orthogonal 16 is built. Right: success vs. B for training sets drawn at random from the 64 · open standalone ↗

The orthogonal 16

One row per instruction. Every value appears 4 times down a column, and every pairing across any two columns appears exactly once — nothing repeated, nothing missing.

object receptacle button

Success rate vs. B

Takeaway: Full coverage is unnecessary — and a well-chosen quarter of it is worth far more than a random quarter.

A rollout is one simulated attempt, scored success or failure. Each point pools all rollouts over 3 training seeds × 3 random draws of the B training instructions; the band is ±1 standard error across those draws (B=4 is a single fixed set, so it has no band). The star is the orthogonal 16 over 3 seeds. The grid is a design diagram, not measured data.

Every task, one at a time

Open full screen ↗

The curves average over instructions; this view breaks that average open. Each small cube is one of the 64, placed by its object, receptacle and button. Pick a training set, then rotate and hover to see how each one did.

Q2: What Is Missing When Sparse Coverage Fails?

Diagonal pretraining learns reusable subtask skills but not instruction steering. Finetuning with only one demonstration per task improves OOD success from 0.4% to 54.7%.

OOD = tasks never seen during pretraining · +1 demo/instruction = one demonstration for each of the 64 task instructions · open standalone ↗

Pretrained only (OOD)
+ 1 demo / instruction (OOD)

OOD Success: Before vs. After 1-Demo Finetuning

L5 Pick & Place (64 tasks) — pretrained on B sparse diagonal tasks, then finetuned with 1 demo per instruction

Takeaway: Even with as few as 4 pretraining tasks, one demonstration per instruction is enough to jump from near-zero to >50% OOD success — the skills were already there.

* B=8 baseline uses a single-seed evaluation (no 3-seed pooled baseline available for that checkpoint).

Q3: What Coverage Is Needed for Dependent Instructions?

Now the two halves of the instruction depend on each other: the robot is told where the second object goes, and must infer where the first one goes.

Dependent Two-Stage Pick-and-Place · 48-task space · every condition trained on B = 16 tasks · open standalone ↗

The task

“Put the cube away, then put the milk on the plate.” Where the cube goes is never stated — but the plate is reserved for the milk, so the cube may go anywhere except 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.
Why “dependent”: the instruction names the milk's container (c2) but never the cube's (c1). Only c1 ≠ c2 pins it down — move the milk and the cube's valid destinations move with it. So training must cover (object, c2) combinations, not objects and containers separately.

What training saw, and what it cost on held-out tasks

A task is an (o1, o2, c2) triple: 4 × 3 × 4 = 48 tasks, spanning 16 (o1, c2) combinations. Every combination in training demonstrates all 3 legal c1 — the rule is always fully shown. Crucially, all four conditions train on the same budget of 16 tasks — what changes is only how widely those 16 are spread over combinations, from 6 up to all 16. Coverage, not data volume.

c2 — reserved for the second object
trained held out c1 = c2, forbidden
16 / 16
combinations trained
48
c1 demonstrations
B = 16
training tasks — same in every condition

measured on the 32 held-out tasks
Task success
± SE

At a fixed 16 training tasks, spreading them wider lifts held-out success 9.2% → 31.4% — no extra data, just better coverage. Seeing the rule in full is not enough; it has to be seen across combinations.

The board's counts and forbidden slots are exact. Which cells a condition covers is illustrative — the logs record how many, not which.

48-task Dependent 2S-PP (l4_2s_no_cont0, 48-task set). Mean ± SE across 5 training seeds, on the 32 held-out tasks. Cov6/9/12 are the B16 runs and Cov16 is the orthogonal 16-task design, so the training budget is 16 tasks in all four conditions — only coverage differs. Training on all 48 tasks is not shown: that run has no held-out split to measure.

Citation

Venue and full bibliographic details to be announced — the entry below will be updated once the paper appears.

@misc{wang2026compositional,
  title  = {Diagnosing Compositional Generalization in Sequential Robot Tasks},
  author = {Wang, Yixiao and Wu, Cheng-En and Sun, Lingfeng and Wang, Pengcheng
            and Ji, Xiang and Liang, Boyuan and Zhan, Guojian and Tomizuka, Masayoshi},
  year   = {2026},
  note   = {Preprint. Venue to be announced}
}