Where Should Action Generation Begin? A Learnable Source Prior for Generative Robot Policies

Resumo Técnico: LeaP - Um Prior de Fonte Aprendível para Políticas de Robôs Generativas

Definição do Problema
Políticas de robôs generativas, incluindo abordagens baseadas em difusão e flow-matching, tipicamente formulam a geração de ações como amostragem condicional. Apesar da diversidade arquitetural, esses métodos compartilham uma limitação comum: eles inicializam a geração de ações a partir de uma distribuição Gaussiana padrão independente de observação, $N(0, I)$. Esta distribuição de origem contém nenhuma informação sobre o estado atual do robô ou o contexto da tarefa. Consequentemente, o modelo generativo deve gastar uma parte significativa de seu "orçamento de integração" para transportar amostras desta fonte de ruído não informada para ações relevantes à tarefa. Embora trabalhos recentes tenham tentado substituir esta Gaussiana padrão por fontes determinísticas derivadas de observações (ex: estimativas de ponto via propriocepção ou latentes visuais), eles negligenciam amplamente a incerteza inerente à distribuição de origem. A questão central abordada por este trabalho é: Onde a geração de ação deve começar? especificamente investigando se a fonte deve ser determinística ou distributiva, e se, sendo distributiva, sua variância deve ser fixa ou adaptativa ao estado.

Metodologia: LeaP
Os autores propõem o LeaP (Learnable source Prior), um framework que substitui a fonte Gaussiana padrão por uma Gaussiana diagonal aprendível e condicionada à propriocepção sobre blocos de ação (action chunks). Os componentes principais são:

1. Rede de Prior ($g_\phi$): Um MLP leve de duas camadas que mapeia características de estado proprioceptivo ($z_s$) para a média ($\mu_\phi$) e a variância adaptativa ao estado ($\sigma^2_\phi$) de uma distribuição Gaussiana diagonal. Crucialmente, o prior condiciona-se apenas à propriocepção, deixando o condicionamento completo (visual + proprioceptivo) para o gerador subsequente.
2. Inicialização Estocástica: Na inferência, a amostra de origem $z_0$ é extraída desta Gaussiana aprendida $q_\phi(z_0 | z_s)$ em vez de $N(0, I)$. Isso fornece uma inicialização que é tanto informada pela observação (ancorada em uma vizinhança dependente do estado) quanto estocástica.
3. Design Plug-and-Play: O LeaP modifica apenas a distribuição de origem. A arquitetura do gerador subsequente (ex: flow-matching ou diffusion-bridge), a dinâmica generativa e o resolvedor de inferência permanecem inalterados.
4. Objetivo de Treinamento Conjunto: A rede de prior e o gerador são treinados de ponta a ponta usando uma função de perda composta:
   • Perda de Flow Matching ($L_{flow}$): Supervisiona o campo de velocidade $v_\theta$ para transportar $z_0$ até a ação alvo.
   • Log-Verossimilhança Negativa ($L_{NLL}$): Supervisiona a densidade condicional do prior, garantindo que a média e a variância previstas alinhem-se com os blocos de ação especialistas.
   • Perda de Alinhamento Contrastivo ($L_{align}$): Uma perda simétrica do tipo CLIP que impõe a correspondência ao nível de instância entre amostras de origem e ações em um espaço de incorporação (embedding) aprendido.

Principais Contribuições

• Reenquadramento da Distribuição de Origem: O artigo identifica a distribuição de origem como um eixo de design independente e reutilizável para políticas de robôs generativas, distinto da escolha das dinâmicas generativas.
• Variância Adaptativa ao Estado: Ao contrário de fontes determinísticas anteriores ou heurísticas de variância fixa, o LeaP aprende conjuntamente a média e uma variância adaptativa ao estado, modelando explicitamente a incerteza da origem.
• Seleção de Modalidade: Os autores demonstram que a propriocepção sozinha é suficiente para parametrizar o prior, e que incorporar características visuais no prior na verdade degrada o desempenho em seu cenário.
• Compatibilidade Plug-and-Play: O método melhora tanto geradores de flow-matching quanto de diffusion-bridge sem alterar suas arquiteturas internas.

Resultados Experimentais
Os autores avaliaram o LeaP em 15 tarefas de manipulação bimanual no ambiente de simulação RoboTwin 2.0 e em três tarefas do mundo real com um robô Franka Research 3.

• Desempenho em Simulação: O LeaP alcançou uma taxa de sucesso média de 81,6%, superando quatro baselines representativas (métodos de fonte determinística A2A e VITA, um contraparte sem prior e uma política de diffusion-bridge) em 6,5 a 25,5 pontos percentuais. A maior lacuna (25,5 pp) foi observada contra o baseline "NoPrior", que compartilhava exatamente a mesma arquitetura de gerador mas utilizava $N(0, I)$, confirmando que o ganho provém do próprio prior aprendível.
• Eficiência: O LeaP utiliza significativamente menos parâmetros (13,22M no total, com o cabeçote do prior respondendo por apenas 0,21M) comparado a baselines como ACT ou BridgePolicy. Ele também exibe uma convergência de treinamento mais rápida e estável.
• Implantação no Mundo Real: Em experimentos no mundo real, o LeaP manteve o desempenho superior com uma taxa de sucesso média de 80,0%, superando o segundo melhor baseline (A2A) em 11,7 pontos percentuais.
• Estudos de Ablação:
  • Modalidade de Entrada: Condicionar o prior apenas à propriocepção rendeu os melhores resultados; fundir características visuais no prior reduziu o desempenho.
  • Parametrização: Aprender conjuntamente a média e a variância adaptativa ao estado foi crítico; alternativas determinísticas ou de variância fixa performaram pior.
  • Geradores: O prior consistentemente melhorou o desempenho através de diferentes backbones de geradores.

Significância e Alegações
O artigo afirma que a distribuição de origem é uma dimensão de design crítica, porém subexplorada, em políticas de robôs generativas. Ao mudar de uma Gaussiana padrão não informada para um prior estocástico aprendível e condicionado à propriocepção, o gerador é liberado do fardo de transportar amostras de uma fonte de ruído distante. Em vez disso, o prior coloca a amostra inicial em uma "região proprioceptivamente plausível", permitindo que o gerador foque no refinamento preciso e condicionado à observação. Os autores afirmam que esta abordagem oferece uma alavanca de baixo custo para melhorar o desempenho da política, a velocidade de convergência e a robustez, complementando as escolhas existentes em dinâmicas generativas.