Componha suas políticas! Melhorando políticas robóticas baseadas em difusão ou fluxo por meio de composição em nível de distribuição durante o teste

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

October 1, 2025

Autores: 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

cs.AI

Resumo

Modelos baseados em difusão para controle robótico, incluindo políticas visão-linguagem-ação (VLA) e visão-ação (VA), demonstraram capacidades significativas. No entanto, seu avanço é limitado pelo alto custo de aquisição de grandes conjuntos de dados de interação. Este trabalho introduz um paradigma alternativo para melhorar o desempenho das políticas sem treinamento adicional do modelo. Surpreendentemente, demonstramos que as políticas compostas podem superar o desempenho de qualquer uma das políticas originais. Nossa contribuição é tripla. Primeiro, estabelecemos uma base teórica mostrando que a composição convexa de escores distribucionais de múltiplos modelos de difusão pode resultar em um objetivo funcional de um passo superior em comparação com qualquer escore individual. Um limite do tipo Gr\"onwall é então usado para mostrar que essa melhoria de um único passo se propaga por trajetórias de geração inteiras, levando a ganhos sistêmicos de desempenho. Segundo, motivados por esses resultados, propomos a Composição Geral de Políticas (GPC), um método sem treinamento que melhora o desempenho combinando os escores distribucionais de múltiplas políticas pré-treinadas por meio de uma combinação convexa e busca em tempo de teste. A GPC é versátil, permitindo a composição plug-and-play de políticas heterogêneas, incluindo modelos VA e VLA, bem como aqueles baseados em difusão ou correspondência de fluxo, independentemente de suas modalidades visuais de entrada. Terceiro, fornecemos uma validação empírica extensa. Experimentos nos benchmarks Robomimic, PushT e RoboTwin, juntamente com avaliações robóticas do mundo real, confirmam que a GPC melhora consistentemente o desempenho e a adaptabilidade em um conjunto diversificado de tarefas. Uma análise adicional de operadores de composição alternativos e estratégias de ponderação oferece insights sobre os mecanismos subjacentes ao sucesso da GPC. Esses resultados estabelecem a GPC como um método simples, porém eficaz, para melhorar o desempenho de controle aproveitando políticas existentes.

English

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\"onwall-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.