Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Compose Your Policies! Improving Diffusion-based or Flow-based Robot Policies via Test-time Distribution-level Composition

Jiahang Cao, Yize Huang, Hanzhong Guo, Rui Zhang, Mu Nan, Weijian Mai, Jiaxu Wang, Hao Cheng, Jingkai Sun, Gang Han, Wen Zhao, Qiang Zhang, Yijie Guo, Qihao Zheng, Chunfeng Song, Xiao Li, Ping Luo, Andrew F. Luo

TL;DR

This work tackles data efficiency and performance plateaus in diffusion- and flow-based robotic policies by proposing General Policy Composition (GPC), a training-free framework that convexly combines distributional scores from multiple pre-trained policies at test time. The authors ground GPC in theory, proving that convex score composition can reduce one-step error and propagate this improvement through trajectory generation via a Grönwall-based stability bound. They formalize a Compositional Diffusion Model (CDM) and show how to integrate heterogeneous VA/VLA policies without retraining, including logical AND/OR operators as extensions. Extensive experiments across Robomimic, PushT, RoboTwin, and real-world tasks demonstrate consistent performance improvements and reveal insights on weight configurations and operator choices, establishing GPC as a simple, versatile method for enhancing robotic control using existing policies.

Abstract

Diffusion-based models for robotic control, including vision-language-action (VLA) and vision-action (VA) policies, have demonstrated significant capabilities. Yet their advancement is constrained by the high cost of acquiring large-scale interaction datasets. This work introduces an alternative paradigm for enhancing policy performance without additional model training. Perhaps surprisingly, we demonstrate that the composed policies can exceed the performance of either parent policy. Our contribution is threefold. First, we establish a theoretical foundation showing that the convex composition of distributional scores from multiple diffusion models can yield a superior one-step functional objective compared to any individual score. A Grönwall-type bound is then used to show that this single-step improvement propagates through entire generation trajectories, leading to systemic performance gains. Second, motivated by these results, we propose General Policy Composition (GPC), a training-free method that enhances performance by combining the distributional scores of multiple pre-trained policies via a convex combination and test-time search. GPC is versatile, allowing for the plug-and-play composition of heterogeneous policies, including VA and VLA models, as well as those based on diffusion or flow-matching, irrespective of their input visual modalities. Third, we provide extensive empirical validation. Experiments on Robomimic, PushT, and RoboTwin benchmarks, alongside real-world robotic evaluations, confirm that GPC consistently improves performance and adaptability across a diverse set of tasks. Further analysis of alternative composition operators and weighting strategies offers insights into the mechanisms underlying the success of GPC. These results establish GPC as a simple yet effective method for improving control performance by leveraging existing policies.

Compose Your Policies! Improving Diffusion-based or Flow-based Robot Policies via Test-time Distribution-level Composition

TL;DR

This work tackles data efficiency and performance plateaus in diffusion- and flow-based robotic policies by proposing General Policy Composition (GPC), a training-free framework that convexly combines distributional scores from multiple pre-trained policies at test time. The authors ground GPC in theory, proving that convex score composition can reduce one-step error and propagate this improvement through trajectory generation via a Grönwall-based stability bound. They formalize a Compositional Diffusion Model (CDM) and show how to integrate heterogeneous VA/VLA policies without retraining, including logical AND/OR operators as extensions. Extensive experiments across Robomimic, PushT, RoboTwin, and real-world tasks demonstrate consistent performance improvements and reveal insights on weight configurations and operator choices, establishing GPC as a simple, versatile method for enhancing robotic control using existing policies.

Abstract

Diffusion-based models for robotic control, including vision-language-action (VLA) and vision-action (VA) policies, have demonstrated significant capabilities. Yet their advancement is constrained by the high cost of acquiring large-scale interaction datasets. This work introduces an alternative paradigm for enhancing policy performance without additional model training. Perhaps surprisingly, we demonstrate that the composed policies can exceed the performance of either parent policy. Our contribution is threefold. First, we establish a theoretical foundation showing that the convex composition of distributional scores from multiple diffusion models can yield a superior one-step functional objective compared to any individual score. A Grönwall-type bound is then used to show that this single-step improvement propagates through entire generation trajectories, leading to systemic performance gains. Second, motivated by these results, we propose General Policy Composition (GPC), a training-free method that enhances performance by combining the distributional scores of multiple pre-trained policies via a convex combination and test-time search. GPC is versatile, allowing for the plug-and-play composition of heterogeneous policies, including VA and VLA models, as well as those based on diffusion or flow-matching, irrespective of their input visual modalities. Third, we provide extensive empirical validation. Experiments on Robomimic, PushT, and RoboTwin benchmarks, alongside real-world robotic evaluations, confirm that GPC consistently improves performance and adaptability across a diverse set of tasks. Further analysis of alternative composition operators and weighting strategies offers insights into the mechanisms underlying the success of GPC. These results establish GPC as a simple yet effective method for improving control performance by leveraging existing policies.

Paper Structure

This paper contains 85 sections, 3 theorems, 80 equations, 29 figures, 20 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Composable Generative Models.
  4. Diffusion Models in Robot Learning.
  5. Compositional Diffusion Models in Robotics.
  6. Preliminaries
  7. Theoretical Analysis of Convex Score Composition
  8. Functional-level improvement
  9. System-level stability
  10. Implications for policy composition
  11. Our Method: General Policy Composition
  12. Compositional Diffusion Models
  13. General Policy Composition
  14. GPC with Superposition
  15. Experiment
  16. ...and 70 more sections

Key Result

Proposition 1

Let two score estimators be $\varepsilon_1=s^*+b_1+\eta_1$ and $\varepsilon_2=s^*+b_2+\eta_2$, with deterministic biases $b_i$ and random zero-mean noise $\eta_i$ that plays the role of the diffusion component in the time-reversed stochastic dynamics (e.g., a reverse-time ODE). For any convex weight with strict inequality whenever the two models’ errors are not perfectly aligned.

Figures (29)

  • Figure 1: Illustration of General Policy Composition.(a) Distributions from pre-trained state-of-the-art diffusion- or flow-based policies can be composed to construct a stronger policy without additional training, with a test-time search over composition weights picking the best parent-policy mix; score composition corresponds to the product of probabilistic density functions (PDFs), steering sampling toward consensus regions. (b) GPC can yield consistent gains across a diverse set of tasks. (c) We find the optimal weight when composing two models can vary depending on the task.
  • Figure 2: Overview of our proposed General Policy Composition. Combining distributional scores from pre-trained diffusion-based or flow-based policies on different conditions (e.g., visual modalities and network backbones), GPC can generate expressive and adaptable action trajectories through convex score combination without additional training.
  • Figure 3: Visualization results of different diffusion policies and the composed policy with GPC. Our proposed GPC can be successful even when one part of the DP fails, and shows better performance when both parts of the DP work.
  • Figure 4: Visual analysis of GPC under different compositions. GPC generalizes across (a) modalities and (b) architectures, with appropriate weighting yielding accurate distributions with better SR than individual policies.
  • Figure 4: Results of GPC with superposition, highlighting performance increase by strong compositional operators.
  • ...and 24 more figures

Theorems & Definitions (3)

  • Proposition 1: Single-step improvement via convex combination
  • Proposition 2: Score-to-sample stability
  • Corollary 1: Convex score combination can reduce overall sampling error