PaperlayerView exact revision on arXiv

Source-linked AI summary

Distillation for Efficient Multitask Manipulation Policies via Conditional Flow Matching

Shreya Deshmukh, Imen Mahdi, Nick Heppert, Abhinav Valada

arXiv:2609.28107v2cs.ROcs.AIcs.LG

TL;DR

Maintaining separate policies becomes costly as task counts grow, while naive joint training can demand more capacity or degrade performance. The paper distills single-task CFM experts into a shared policy and combines their guidance with the original CFM objective. On RLBench, both variants outperform naive multi-task training and the average single-task expert, with the full objective reaching 79.5% average success rate.

  • Problem

    Separate policies become increasingly impractical as task counts grow, while naive joint training can require greater capacity or degrade per-task performance.

    Source: §I · p. 1, B. Multi-Task policy distillation · p. 2

  • Method

    The method distills single-task CFM expert velocity fields or generated target samples into one multi-task policy while retaining the original CFM objective.

    Source: §I · p. 1

  • Results

    Both variants outperform a naive multi-task baseline and the average of single-task experts, while the full objective achieves 79.5% average success rate.

    Source: D. Evaluation Results · p. 4, §VI · p. 5

  • Takeaways & Limitations

    The reported benefit of distillation grows as more tasks are added until the student saturates its model capacity.

    Source: §VI · p. 5

    Evaluation is limited to simulated RLBench tasks with single-arm manipulation; validation on more benchmarks and real robots remains future work.

    Source: §V · p. 5

Abstract

from arXiv · show

Advances in generative modeling have recently been extensively employed in robotics for policy learning. In particular, Conditional Flow Matching (CFM) trained with expert demonstrations has been shown to outperform existing methods on robot manipulation benchmarks. While prior work has mainly focused on single-task settings, we study the problem from a multi-task perspective, as training independent models for each task is computationally expensive. Multi-Task policy learning comes with its own set of challenges, as naively training on a concatenated dataset of demonstrations would either require increased model capacity to accommodate the added complexity or result in drops in performance. We propose to distill knowledge from single-task CFM experts into a shared multi-task policy by transferring their learned velocity fields. We combine this distillation signal with the original CFM objective to retain fidelity to the demonstrations. Experiments on RLBench show that our approach improves multi-task policy performance over naive training while maintaining a fixed model size.

I. INTRODUCTION

The paper addresses the cost of maintaining a separate policy for every task and the negative transfer that can arise when tasks are jointly trained under fixed model capacity. It distills single-task CFM experts into a shared policy using expert velocity fields or generated target samples, alongside the original CFM objective.

I. INTRODUCTION

  • Separate task policies become increasingly impractical as task counts grow, while fixed-capacity joint training faces negative transfer.Large vision-language-action models also make real-time deployment and task-specific finetuning computationally expensive.Source: §I · p. 1
  • CFM learns a velocity field for generative policies, but a single model trained naively across tasks can encounter conflicting gradients.The supplied passage describes this challenge but does not complete its explanation.Source: §I · p. 1
  • The paper studies two distillation variants using a toy example and robotic manipulation experiments.The variants use expert velocity fields and expert-generated target samples.Source: §I · p. 1
  • The proposed approach distills single-task CFM experts into one multi-task policy without increasing student model size.It transfers either expert velocity fields or expert-generated target samples and retains the original CFM objective.Source: §I · p. 1

II. PREREQUISITES: CONDITIONAL FLOW MATCHING

Conditional Flow Matching models generation as a continuous transformation from a simple initial distribution toward the data distribution. It trains a neural velocity field to match conditional target velocities at interpolated samples, then integrates that field to generate samples.

II. PREREQUISITES: CONDITIONAL FLOW MATCHING

  • CFM uses a time-dependent vector field to transform samples from a simple initial distribution into samples from the data distribution.The field induces a flow defined as the solution to an ordinary differential equation.Source: §II · p. 1
  • A linear conditional path interpolates independently drawn initial and target samples as z_t = (1 − t) z_0 + t z_1.The path is defined for t ∈ [0, 1].Source: §II · p. 1
  • Training regresses the neural velocity field against conditional ground-truth velocities at constructed intermediate samples; inference integrates the learned field.The provided method overview also depicts single-task CFM experts distilled into a multi-task model, with aggregated ground-truth data anchoring final samples.Source: §II · p. 2, Figure 1 · p. 2

III. MULTI-TASK POLICY DISTILLATION WITH CONDITIONAL FLOW MATCHING

The formulation learns one policy for a set of tasks from expert demonstrations pooled across task-specific datasets. It casts action trajectories or chunks as the flow variable and conditions the learned velocity field on the paired state.

A. Problem Formulation

  • The learning setting seeks one policy capable of solving a set of K tasks using expert demonstrations collected across those tasks.The total demonstration count is the sum of the per-task dataset sizes.Source: A. Problem Formulation · p. 2
  • The union of the task-specific datasets is denoted D.Each task has its own dataset D_k within the pooled demonstrations.Source: A. Problem Formulation · p. 2
  • Each demonstration pairs a state with an action trajectory or chunk, which serves as the flow target while the velocity field is conditioned on the state.The target z_1 is the demonstrated action a.Source: A. Problem Formulation · p. 2

B. Multi-Task policy distillation

The method transfers knowledge from task-specific CFM experts because their velocity fields can provide smoother guidance than raw targets in multi-task learning. It combines expert distillation with the multi-task CFM loss to preserve fidelity to the demonstrations.

B. Multi-Task policy distillation

  • Naively training on pooled task data can require greater model capacity or degrade per-task performance because conditional velocity targets vary across tasks and demonstrations.The paper motivates explicit transfer by noting overlap in operational space and trajectory behavior alongside task-specific differences.Source: B. Multi-Task policy distillation · p. 2
  • The method trains single-task CFM experts, then uses the expert associated with each sampled demonstration's task to supervise the multi-task model at the same interpolated point.The expert prediction is described as a smoother, less noisy target than the ground-truth velocity, particularly early in integration.Source: B. Multi-Task policy distillation · p. 2
  • The multi-task model combines expert distillation with its CFM loss to remain anchored to the original data distribution.A fixed scalar λ weights the distillation term relative to the ground-truth CFM loss.Source: B. Multi-Task policy distillation · pp. 2–3

C. Distilling Learned Target Distribution vs. Velocity Vector Fields

The learned single-task experts generate synthetic action targets that augment demonstrations, while velocity-field distillation transfers expert guidance. Expert-based augmentation can broaden target coverage, but its effectiveness depends on expert quality.

C. Distilling Learned Target Distribution vs. Velocity Vector Fields

  • Expert-generated targets can increase target-distribution coverage beyond the fixed demonstrations without additional data collection.The experts synthesize plausible actions from their learned task distributions.Source: C. Distilling Learned Target… · p. 3
  • For each demonstration state, the method integrates a frozen single-task expert from sampled noise to generate synthetic action targets.Repeating the process with multiple noise samples adds synthetic state-action pairs to the original demonstrations.Source: C. Distilling Learned Target… · p. 3
  • Because generated actions are learned samples rather than ground truth, expert modeling errors can transfer to student velocity targets.The effectiveness of this augmentation therefore depends strongly on the quality of the single-task experts.Source: C. Distilling Learned Target… · p. 3

IV. EXPERIMENTAL EVALUATION

The evaluation asks whether expert distillation helps multi-task manipulation policies and how its variants behave. It also examines performance as task count and distillation weight change.

IV. EXPERIMENTAL EVALUATION

  • The evaluation examines whether expert distillation enables positive cross-task transfer on a real manipulation benchmark compared with naive multi-task training.It also asks what intuition explains the benefit and what limitations and benefits each variant has.Source: §IV · p. 3
  • The study compares velocity-field distillation with expert-generated target-sample distillation in balancing knowledge transfer against fidelity to true action distributions.Source: §IV · p. 3
  • The evaluation asks how distillation benefits change as the number of tasks grows and how performance responds to the distillation weight λ.Source: §IV · p. 3

A. Toy Example: 1D Gaussian Distributions

In a two-mode Gaussian toy example, the combined objective most closely matches both target distributions. The alternatives illustrate trade-offs between shared velocity guidance, target matching, and inherited expert errors.

A. Toy Example: 1D Gaussian Distributions

  • The combined objective, which distills expert velocity fields while anchoring the flow to real data, produces samples that most closely match the target distributions.Source: A. Toy Example: 1D Gaussian… · p. 3
  • Naive multi-task training shares early-flow velocity fields across tasks but yields a higher-variance final distribution.The passage attributes this to a noisier learning signal when jointly optimizing for both tasks.Source: A. Toy Example: 1D Gaussian… · p. 3
  • Velocity-field-only distillation shares early-stage behavior, but without direct target-distribution matching its final samples do not recover either mode fully.Source: A. Toy Example: 1D Gaussian… · p. 3
  • Training on expert-generated samples yields final distributions closer to the targets, but its velocity field is sensitive to inherited expert errors and variance.Source: A. Toy Example: 1D Gaussian… · p. 3

B. Benchmark

The benchmark evaluates policies on eight RLBench manipulation tasks, comparing single-task, naive multi-task, amalgamation, and two distillation variants. The setup uses shared model architecture and observation conditioning without an explicit task identifier.

B. Benchmark

  • Evaluation uses eight RLBench manipulation tasks, 100 expert trajectories per task, and 100 randomized evaluation episodes per task, reporting per-task and mean success rates.Source: B. Benchmark · p. 3

C. Baselines

  • The baselines include task-specific experts and a same-architecture naive multi-task policy trained on the combined task datasets.The multi-task policy receives no task identifier; the setup assumes task identity can be inferred from the state.Source: C. Baselines · p. 3
  • The comparison includes knowledge amalgamation, which aligns task-specific feature spaces with a probing layer, and both proposed distillation variants.Source: C. Baselines · pp. 3–4
  • In the toy comparison, the combined method most closely matches both target modes, whereas the alternatives show higher variance, incomplete mode recovery, or inherited expert noise.Source: Figure 2 · p. 4
  • Table I reports per-task success rates across eight RLBench tasks for the single-task, naive multi-task, amalgamation, and two proposed variants.The caption says the best result per task is bolded but supplies no values.Source: Table I · p. 4
  • The shared policy uses PointFlowMatch with point-cloud and proprioceptive features conditioning a 1D U-Net, and inference numerically integrates its learned ODE.The multi-task model uses the same architecture as the single-task experts and receives observations without an explicit task label.Source: C. Baselines · p. 4

D. Evaluation Results

The full distillation objective performs best overall, combining expert velocity fields with an anchor to real demonstrations. Both method variants outperform the average single-task expert, with the largest stated task-level gain on open fridge.

D. Evaluation Results

  • Knowledge Amalgamation exceeds the single-task average, while trajectory distillation improves over naive multi-task training but remains limited by expert-policy noise.The naive aggregated-data policy underperforms the average single-task expert; expert success rates range from 35% to 99%.Source: D. Evaluation Results · p. 4
  • 79.5% average success rate: the full objective achieves the best overall performance and both variants improve over the average single-task expert.The objective combines expert velocity fields with an anchoring term that encourages the flow toward real data points.Source: D. Evaluation Results · pp. 4–5
  • Open fridge success rises from 46% to 77%, with the task benefiting from knowledge from the similar open box task.The authors identify this as most apparent among the least-performing tasks.Source: D. Evaluation Results · p. 5

E. Ablations

Across task counts, matching expert velocity fields while anchoring to real data outperforms naive multi-task training and, through eight tasks, the single-task expert average. The sensitivity study favors balancing distillation and ground-truth CFM losses; at sixteen tasks, a larger student recovers the positive trend.

a) Scalability

  • The naive multi-task baseline underperforms the single-task expert average across task counts, even when its model capacity is increased at K = 16.Simply distilling expert-generated trajectories also does not reliably improve performance because it is sensitive to errors in those trajectories.Source: E. Ablations · p. 5, Figure 3 · p. 5
  • At K = 16, performance falls to approximately the single-task expert average, but a larger student recovers the positive trend.This comparison distinguishes a student-capacity limit from a fundamental limitation of the distillation approach.Source: E. Ablations · p. 5, Figure 3 · p. 5
  • At K = 8, the method reaches 79.5% SR versus 64.4% for the single-task expert average.Across K ≤ 8, the combined objective also outperforms that average and consistently improves over the naive baseline.Source: Figure 3 · p. 5, E. Ablations · p. 5

b) Sensitivity

  • λ = 0.5 yields the best average success rate, 79.5%, while λ = 0 gives 61.4% and λ = 0.25, 0.75, and 1.0 give 69.0%, 68.0%, and 65.4%, respectively.The sweep covers average success rate across eight RLBench tasks; λ = 0 is the naive multi-task objective.Source: E. Ablations · p. 5, Table II · p. 5

V. LIMITATIONS AND FUTURE WORK

The evaluation is limited to simulated RLBench single-arm manipulation, and the distillation weight is fixed. The conclusion reports that both distillation variants outperform naive multi-task training and the average single-task expert until student capacity saturates.

V. LIMITATIONS AND FUTURE WORK

  • Evaluation is limited to simulated RLBench tasks with single-arm manipulation; validation on more benchmarks and real robots remains future work.The study also uses a fixed distillation weight λ and proposes investigating a time-dependent schedule.Source: §V · p. 5

VI. CONCLUSION

  • Both variants outperform naive multi-task training and the average single-task experts, with benefits growing as tasks are added until the student saturates its capacity.The shared compact policy is trained using expert velocity fields or expert-augmented target actions.Source: §VI · p. 5
Loading 2609.28107v2…