Source-linked AI summary

Task-Prototype Guided Flow Matching for Few-Shot Generalization in Vision-Language Robot Manipulation

Yizhao Wang, Guantao Zhang, Jingbo Wang

arXiv:2609.27780v1cs.ROcs.CV

TL;DR

Language and observations can underspecify manipulation execution, making adaptation from a few demonstrations difficult. TP-Flow extracts structured task prototypes from support demonstrations and uses them to guide the flow prior and velocity field. On LEROBOT-ARM-SO101, it improves few-shot and held-out-task generalization while maintaining real-time closed-loop execution.

  • Problem

    Language and current observations may underspecify execution details such as contact timing and corrective motion, while few-shot support sets can contain nuisance information.

    Source: §1 · pp. 1–2
  • Method

    TP-Flow extracts structured task-prototype tokens from few-shot demonstrations and uses them to guide the initial flow prior and gated velocity-field modulation.

    Source: §6 · p. 26
  • Results

    On LEROBOT-ARM-SO101, TP-Flow improves few-shot generalization across held-out target tasks and maintains real-time closed-loop execution.

    Source: §6 · p. 26
  • Takeaways & Limitations

    Prototype distance aligns with action-distribution distance, supporting task prototypes as measurable representations of task-level execution structure.

    Source: §6 · p. 26

    Evaluation is presently limited to LEROBOT-ARM-SO101; cross-embodiment evaluation remains a next step.

    Source: §5 · p. 26

Abstract

from arXiv · show

Vision-language robot manipulation policies can follow semantic instructions, but adapting them to a new procedure from only a few demonstrations remains difficult because language underspecifies contact timing, motion phases, corrective behavior, and execution style. This paper presents Task-Prototype Guided Flow Matching (TP-Flow), a few-shot manipulation framework that converts support demonstrations into structured task-prototype tokens and uses them to guide both the initial flow prior and the velocity field. TP-Flow employs symmetric cross-attention with learnable queries to extract phase-level prototypes, parameterizes a task-adaptive initial distribution, and injects prototype information through gated adaptive normalization. It is trained with an episodic support-query objective and prototype contrastive regularization, so few-shot adaptation is simulated during training while nuisance information is suppressed. On the LEROBOT-ARM-SO101 platform, TP-Flow achieves 66.8\%, 79.6\%, and 82.1\% success rates under 1-, 4-, and 6-shot settings, with a 75.5\% few-shot AUC. At 1-shot, it improves over CFM, Pooled-Demo CFM, and In-Context Flow by 29.8, 14.5, and 9.9 percentage points. It also improves held-out target-group generalization across novel-object transfer, goal recombination, long-horizon composition, and contact/correction tasks. TP-Flow maintains real-time execution with six online prototype tokens, 54.3 ms latency, 3.9 GB peak memory, and a 10 Hz control rate, while reducing the noisy-support success drop to 6.2\%. Theoretical diagnostics show that prototype distance aligns with action-distribution distance, the adaptive prior reduces transport cost, and gated modulation keeps measured trajectory deviations below the derived ODE bound. The code repository is omitted for anonymous review.

1. Introduction

Language and current observations do not fully specify how a manipulation procedure should be executed, while a few demonstrations contain both useful execution cues and nuisance variation. TP-Flow turns those demonstrations into structured prototypes that guide action generation.

1. Introduction

  • Few-shot adaptation is difficult because language underspecifies contact handling and motion phases, while support demonstrations may mix execution cues with nuisance information.The paper frames the question as how to represent demonstrations and let that representation influence action generation.Source: §1 · p. 2
  • Average pooling can discard temporal structure, raw in-context demonstrations raise online cost, and feature concatenation leaves support influence on action generation unspecified.Over-conditioning can inject support noise, whereas under-conditioning can collapse to language-only control.Source: §1 · p. 2
  • TP-Flow uses symmetric cross-attention to extract phase-level prototypes that guide a task-adaptive initial flow prior and modulate the velocity field through gated adaptive normalization.Episodic support-query training and prototype contrastive regularization simulate few-shot adaptation and suppress nuisance variation.Source: §1 · p. 2
  • The evaluation covers few-shot adaptation, cross-task generalization, novel-object robustness, long-horizon transfer, inference efficiency, and module-level ablations on LEROBOT-ARM-SO101.The introduction motivates generative action policies by their ability to represent multiple valid trajectories, while noting that ordinary conditional flow matching remains under-specified for few-shot cross-task adaptation.Source: §1 · pp. 2–3

2. Related Work

Prior work supplies scalable vision-language policies, few-shot benchmarks, generative action models, and task representations, but leaves open how a small demonstration set should shape action-generation dynamics. TP-Flow addresses that gap with structured prototype tokens.

2.1 Vision-language-action models for robot manipulation

  • VLA systems scale multimodal policies and robot data, but generally lack an explicit few-shot execution prototype; TP-Flow adds structured conditioning without requiring a new foundation backbone.Prior work also grounds language in spatial and executable motion structure, whereas TP-Flow extracts that structure from few-shot demonstrations.Source: §2.1 · p. 3

2.2 Few-shot robot adaptation

  • Few-shot benchmarks evaluate adaptation to new goals, scenes, and task compositions, while prior representation and imitation methods do not directly specify how support data shapes a generative trajectory at test time.TP-Flow uses episodic support-query training to match the few-shot evaluation setting.Source: §2.2 · p. 4

2.3 Generative policies and flow matching

  • Generative policies model multiple valid trajectories, and flow-based robot policies support continuous control, but most treat task context as a standard condition rather than a mechanism shaping prior and velocity.TP-Flow instead uses structured observation-action prototype tokens to control both.Source: §2.3 · pp. 4–5

2.4 Task representation for cross-task generalization

  • Language, pooled demonstrations, in-context prompts, and learned embeddings each leave open how few-shot demonstrations should control action-generation dynamics; TP-Flow uses multiple prototypes to represent phases and guide the flow prior and gated velocity modulation.Its contrastive objective encourages same-task demonstrations to remain close despite scene variation while separating different task dynamics.Source: §2.4 · p. 5

3. Method

TP-Flow models action generation as flow matching conditioned on current observations and a few-shot support set. It encodes support demonstrations into fixed-size prototypes that guide both the initial prior and the velocity field.

3.1 Problem setup and flow objective

  • The policy receives RGB, proprioceptive state, and language, predicts an action chunk, and uses a small support set containing execution cues such as approach direction, contact timing, correction behavior, and release phase.A support-query episodic protocol maps the support set to prototypes and trains on query trajectories; flow matching transports a prototype-conditioned initial action sample toward an expert action chunk.Source: §3.1 · p. 6

3.2 Overview

  • TP-Flow converts multimodal query and support frames into observation-action tokens, compresses support into fixed-size prototypes, then uses them to parameterize the initial flow prior and modulate the velocity field.Unlike a single averaged feature, its multiple prototypes represent distinct manipulation phases, and training uses support-query episodes.Source: §3.2 · pp. 6–7, Figure 1 · p. 7

3.3 Multimodal observation encoding

  • A shared multimodal encoder fuses RGB, robot state, and language into the current condition and reuses its weights for support frames; support-specific action-temporal fusion captures gripper timing and local motion tendency.The support observation-action tokens are stacked as input to the prototype encoder.Source: §3.3 · pp. 7–8, Figure 2 · p. 8

3.4 Structured prototype learning with symmetric cross-attention

The prototype encoder compresses variable-length support tokens into a fixed-size set whose learnable queries specialize to reusable manipulation phases. Symmetric cross-attention lets prototypes retrieve support information and support tokens attend back to the prototypes.

3.4 Structured prototype learning with symmetric cross-attention

  • The encoder maps variable-length support tokens to M prototype tokens, with M much smaller than the support length while preserving phase-level execution cues.This fixed-size set avoids average pooling so different execution phases can influence action generation separately.Source: §3.4 · pp. 8–9
  • Learnable queries are not assigned fixed semantic labels; during training they specialize to reusable manipulation phases such as approach, grasp, contact, correction, transport, and release.A query can attend to support frames associated with gripper closure or corrective motion after contact.Source: §3.4 · p. 9
  • Bidirectional attention lets prototype queries retrieve support information and support tokens attend back to the prototypes, checking whether a prototype explains task-relevant support content.Background, lighting, and accidental visual details are less likely to remain stable under this interaction.Source: §3.4 · p. 9

3.5 Prototype-guided flow prior

TP-Flow uses support-derived prototypes to adapt the initial action distribution, since manipulation tasks can require different degrees of trajectory diversity. The adaptive prior changes the flow's starting point, allowing support demonstrations to shape the full generation trajectory.

3.5 Prototype-guided flow prior

  • Support-derived prototypes guide the initial action prior, which can be narrower for precise insertion or broader for placing.This replaces the standard isotropic Gaussian starting assumption.Source: §3.5 · pp. 9–10
  • A pooled summary of the prototype tokens represents overall task execution style and feeds the prior network.Pooling may use mean or attention pooling.Source: §3.5 · p. 9
  • The prior network predicts a mean that shifts the starting distribution toward a task-specific action manifold.The standard deviation is bounded to avoid deterministic or extremely noisy initial samples.Source: §3.5 · pp. 9–10
  • The initial flow sample is drawn by reparameterizing the task-adaptive distribution, so support demonstrations affect the whole flow trajectory.The sampling equation uses element-wise multiplication.Source: §3.5 · p. 10

3.6 Gated prototype modulation of the velocity field

TP-Flow conditions its velocity field on the current action sample, flow time, observation condition, and task prototypes. Cross-attention turns the prototypes into gated layer-wise modulation rather than concatenating raw support tokens to hidden features.

3.6 Gated prototype modulation of the velocity field

  • At flow time 0 the sample comes from the prototype-guided prior, and at flow time 1 it reaches the expert action chunk.Inference generates actions by numerically integrating the learned velocity field.Source: §3.6 · pp. 10, 12
  • The network embeds the current flow state, flow time, and observation condition, then uses the hidden state to query the prototype tokens.The velocity field is conditioned on these inputs and the few-shot execution context.Source: §3.6 · p. 10
  • Each velocity layer queries task prototypes to generate scale, shift, and gate parameters for gated adaptive normalization.The gate can softly attenuate prototype information when support is noisy or visually inconsistent.Source: §3.6 · p. 10, Figure 3 · p. 11
  • Gated adaptive normalization lets prototypes alter internal computation without forcing raw support tokens into the velocity network.A gated residual path controls how much modulated information enters the hidden state.Source: §3.6 · p. 12
  • The layer-wise mechanism is more flexible than hard concatenation and more explicit than language-only conditioning.Source: §3.6 · p. 12

3.7 Episodic training and prototype contrastive regularization

TP-Flow simulates few-shot adaptation during training by dividing each sampled task's demonstrations into support and query sets. Its regularizers encourage task-relevant prototypes, stable priors, and smooth velocity changes.

3.7 Episodic training and prototype contrastive regularization

  • Each training episode splits a sampled task's demonstrations into support and query sets, using the former for prototypes and the latter for flow-matching supervision.This simulates few-shot adaptation during training rather than adding demonstrations only at test time.Source: §3.7 · p. 12
  • Two support views from the same task are sampled to reduce overfitting to background, lighting, or accidental object poses.Source: §3.7 · p. 12
  • Contrastive regularization pulls support encodings of the same task together and separates prototypes from different tasks, encouraging task structure over incidental visual details.Source: §3.7 · p. 13
  • The prior KL term keeps the prototype-guided prior near a well-behaved Gaussian unless support evidence favors a strong task-specific shift.Source: §3.7 · p. 13
  • Velocity-smoothness regularization discourages abrupt changes along the generation trajectory, while local prototype consistency aligns prototype and empirical action-distribution neighborhoods.The full objective balances these terms against flow matching, prototype invariance, and prior stability.Source: §3.7 · p. 13

3.8 Training and inference

Algorithm 1 summarizes TP-Flow's episodic training procedure, and inference reuses support-derived prototypes while generating actions through the learned flow. The method's reported theoretical diagnostics examine prototype consistency, prior transport cost, and gated-modulation stability.

3.8 Training and inference

  • Algorithm 1 combines support-query sampling, prototype extraction, adaptive-prior estimation, flow-matching supervision, and prototype, consistency, and smoothness regularization.Source: §3.8 · p. 13
  • At test time, TP-Flow computes prototypes once from K demonstrations and reuses them for subsequent query observations.The robot executes only the first action or a short chunk prefix before replanning; inference cost depends on M prototype tokens rather than full support trajectory length.Source: §3.8 · p. 14
  • Supporting-material diagnostics examine prototype/action-distance correlation, adaptive-prior transport cost, and gated ODE perturbations on LEROBOT-ARM-SO101.The accompanying analysis concerns local prototype consistency, prior variance reduction, and stability under noisy support demonstrations.Source: §3.8 · p. 14

4. Experiments

Experiments on LEROBOT-ARM-SO101 assess few-shot success, held-out task generalization, action quality, robustness, and inference efficiency.

4.1 Experimental protocol and task design

  • Few-shot evaluation freezes source-task-trained weights, uses support demonstrations only to compute a task prototype, and tests on new query states without test-time gradient updates.The task set includes held-out target groups and seven representative executable tasks.Source: §4.1 · pp. 14–15

4.2 Evaluation metrics and implementation details

  • Evaluation measures task success and completion horizon, few-shot gains and FS-AUC, motion quality, target-group generalization, support-noise robustness, and inference cost.Implementation uses fixed learnable prototype queries, so online velocity decoding attends to prototype tokens rather than the full support sequence.Source: §4.2 · pp. 16–17

4.3 Main results on few-shot generalization

  • At 1-shot, TP-Flow reaches 66.8% SR, improving over CFM by 29.8 points, Pooled-Demo CFM by 14.5 points, and In-Context Flow by 9.9 points; at 6-shot it reaches 82.1% SR and 75.5% FS-AUC.The study reports consistent gains across four support sizes.Source: §4.3 · p. 17, Table 3 · p. 18

4.4 Cross-task and novel-object generalization

  • TP-Flow achieves the highest success rate across spatial, object, goal, and long-horizon target groups, with gains largest for robustness and contact/correction settings.Long-horizon errors still accumulate from early alignment, pickup, or reorientation mistakes; ambiguous contact and dynamically incompatible corrections remain failure conditions.Source: Table 4 · p. 18, §4.3 · p. 18, §4.4 · pp. 20, 22

4.3 Main results on few-shot generalization

  • At 4-shot, TP-Flow reduces smoothness error to 0.052 and replanning error to 0.031, raises contact recovery to 72.8%, and limits noisy-support success drop to 6.2.It also has the lowest action norm and fewest jerk events among compared methods.Source: §4.3 · p. 18, Table 5 · p. 19, Table 6 · p. 20

4.5 Ablation study

  • Replacing the structured prototype encoder lowers SR@4 from 79.6% to 68.8%, while removing contrastive regularization lowers prototype stability from 0.91 to 0.76.Removing the adaptive prior lowers SR@4 to 73.5%, and removing gated modulation raises noisy-support success drop to 14.4.Source: Table 8 · p. 22, §4.5 · pp. 22–23
  • Six prototype tokens and eight ODE integration steps are selected as practical operating points balancing success and latency.More tokens improve SR@4 until success saturates beyond six; gains from additional solver steps become marginal after eight.Source: Figure 10 · p. 23, §4.5 · p. 23

4.6 Inference efficiency

  • Under 4-shot evaluation, TP-Flow uses six structured prototype tokens online, reaches 54.3 ms latency, and maintains 10 Hz control; In-Context Flow uses 64-token context, 91.4 ms latency, and 6.8 GB peak memory.TP-Flow’s online latency remains nearly flat as the support set grows because it is compressed before action generation.Source: Table 9 · p. 23, §4.6 · p. 24

5. Discussion

TP-Flow treats demonstrations as an execution-level task representation and uses compact prototypes to guide action generation. The discussion connects this design to adaptation, robustness, efficiency, and the method’s current scope limits.

Why task prototypes matter

  • Structured prototypes capture execution details such as approach, contact timing, recovery, and phase ordering that language alone may not specify.On LEROBOT-ARM-SO101, structured tokens outperform language-only conditioning, pooled demonstration features, and raw in-context support tokens.Source: §5 · p. 24

Prototype geometry

  • Prototype distance tracks empirical action-distribution distance, and additional demonstrations improve correlation and same-group retrieval while refining compact prototype neighborhoods.The embedding map places tasks with similar manipulation structure nearby rather than collapsing them into one language-only cluster.Source: §5 · pp. 24–25

Prototype guidance beyond feature conditioning

  • Prototypes both move initial samples toward task-specific action manifolds through an adaptive prior and control velocity-field modulation through gated AdaLN.Supporting-material diagnostics report lower transport cost and loss variance from the adaptive prior, and lower sensitivity to noisy prototype perturbations from gated modulation.Source: §5 · p. 25

Few-shot generalization through episodic adaptation

  • Episodic support-query training makes few-shot adaptation part of the learned objective, with most success-curve gains occurring from one to four demonstrations.The smaller improvement from four to six demonstrations suggests that additional examples mainly refine a representation once the main action phases are covered.Source: §5 · p. 25

Robustness to noisy support demonstrations

  • Prototype contrastive regularization and gated AdaLN address noisy support by compacting same-task representations and allowing the velocity network to reduce unreliable prototype influence.Under increasing support perturbation, gate activation decreases and trajectory error grows more slowly than with direct prototype concatenation.Source: §5 · p. 25

Efficiency and limitations

  • TP-Flow encodes the support set once into a small, fixed-size set of prototype tokens, keeping its policy within the 10 Hz control target while preserving much of demonstration conditioning’s adaptation benefit.A single pooled vector loses temporal and phase-level structure. Wrong-goal, incomplete, or misleading demonstrations can bias the generated trajectory.Source: §5 · p. 26

6. Conclusion

TP-Flow converts few-shot support demonstrations into structured prototypes that guide both the initial flow prior and velocity field. On LEROBOT-ARM-SO101, it improves generalization across several held-out settings, while diagnostics characterize prototype geometry and flow behavior.

6. Conclusion

  • TP-Flow extracts phase-level execution structure through symmetric cross-attention, adapts the initial prior, and modulates the velocity network with gated adaptive normalization.The support demonstrations become structured task-prototype tokens rather than a pooled condition or long online context.Source: §6 · p. 26
  • Experiments on LEROBOT-ARM-SO101 show improved few-shot generalization across held-out target tasks, novel objects, goal recombination, long-horizon execution, and noisy-support settings.The paper reports stronger success rates than behavior-cloning, VLA-style, action-chunking, diffusion-policy, and non-prototype flow-matching baselines, with real-time closed-loop execution.Source: §6 · p. 26
  • Diagnostics show prototype distances align with action-distribution distances, the adaptive prior reduces transport cost and loss variance, and gated modulation keeps measured deviations below the derived ODE bound.The authors identify cross-embodiment transfer and tactile- or force-aware manipulation as future work.Source: §6 · pp. 26–27

Declarations

The declarations report no conflicts of interest and mark ethics approval and consent to participate as not applicable; authorship contributions are omitted for anonymous review.

Declarations

Loading 2609.27780v1…