Recovering Hidden Reward in Diffusion-Based Policies

1. Declaração do Problema

Políticas baseadas em difusão tornaram-se um paradigma de última geração para aprender habilidades de manipulação robótica a partir de demonstrações de especialistas, devido à sua capacidade de modelar distribuições de ação complexas e multimodais. No entanto, essas políticas geralmente dependem de Clonagem de Comportamento (BC), que aprende a imitar ações sem modelar explicitamente a intenção subjacente ou a estrutura de recompensa.

• Limitações da Difusão Padrão: Elas carecem de uma representação explícita de energia, tornando-as propensas a falhar quando situações no momento do teste se desviam da distribuição de treinamento (Fora da Distribuição ou OOD). Elas não podem fornecer facilmente um sinal de recompensa para Aprendizado por Reforço (RL) a jusante sem recorrer a treinamento adversarial instável ou Aprendizado por Reforço Inverso (IRL) computacionalmente caro.
• A Lacuna: Há necessidade de um framework que unifique a modelagem generativa de ações com a recuperação de recompensa, permitindo tanto imitação de alta qualidade quanto a extração de um sinal de recompensa válido para refinamento da política, garantindo ao mesmo tempo garantias teóricas sobre generalização.

2. Metodologia: ENERGYFLOW

Os autores propõem ENERGYFLOW, um framework que parametriza uma função de energia escalar $E_\theta(s, a)$ cujo gradiente serve como campo de remoção de ruído para geração de ações.

Arquitetura Central

• Parametrização de Energia Escalar: Ao contrário das políticas de difusão padrão que regressam diretamente uma previsão de ruído vetorial $\epsilon_\theta(s, a)$, o ENERGYFLOW aprende uma função de energia escalar $E_\theta(s, a, t)$.
• Campo Vetorial Conservativo: A função de pontuação (campo de remoção de ruído) é derivada via diferenciação automática:
  $$S_\theta(a, s, t) = -\nabla_a E_\theta(a, s, t)$$
  Por construção, isso garante que o campo vetorial seja conservativo (sem rotacional, $\nabla \times S = 0$), significando que é o gradiente de um potencial escalar. Isso impõe integrabilidade, garantindo que as preferências de ação sejam transitivas e fisicamente realizáveis.
• Treinamento: O modelo é treinado usando Correspondência de Pontuação de Remoção de Ruído (DSM). O objetivo minimiza a diferença entre o gradiente aprendido e a direção real do ruído:
  $$\mathcal{L} = \mathbb{E}_{t, a_0, \epsilon} \left[ \sigma^2(t) \left\| -\nabla_{a_t} E_\theta(a_t, s, t) + \frac{\epsilon}{\sigma(t)} \right\|^2 \right]$$
• Inferência: As ações são geradas integrando a EDO de fluxo de probabilidade usando o gradiente da função de energia aprendida.

Extração de Recompensa

O framework extrai um sinal de recompensa diretamente da paisagem de energia aprendida. Para lidar com a constante de integração dependente do estado $c(s)$ (que torna os valores brutos de energia ambíguos entre diferentes estados), os autores propõem Moldagem Centralizada:
$$\tilde{r}(a, s) = -\left( E_\theta(a, s) - \mathbb{E}_{a' \sim \mathcal{N}}[E_\theta(a', s)] \right)$$
Isso subtrai uma linha de base dependente do estado (estimada via amostragem de Monte Carlo), normalizando a recompensa para refletir apenas preferências de ação relativas dentro de um estado.

3. Contribuições Teóricas Chave

O artigo fornece provas teóricas rigorosas estabelecendo a conexão entre correspondência de pontuação de difusão e recuperação de recompensa:

1. Equivalência Pontuação-Recompensa: Sob a suposição de Otimização de Entropia Máxima, o artigo prova que a função de pontuação aprendida por modelos de difusão é proporcional ao gradiente da função Q suave do especialista ($\nabla_a Q^* \propto S^*$). Assim, a correspondência de pontuação recupera implicitamente o gradiente de recompensa sem otimização min-max adversarial.
2. Redução de Complexidade via Restrições Conservativas: Os autores provam que restringir o espaço de hipóteses a campos conservativos (gradientes de escalares) reduz estritamente a complexidade de Rademacher em comparação com campos vetoriais não restritos. Especificamente, o limite de complexidade para campos conservativos escala com a suavidade da representação ($L$), enquanto campos não restritos escalam com a dimensão de saída ($\sqrt{d}$).
3. Limites de Generalização OOD: Aproveitando a complexidade reduzida, o artigo deriva limites de generalização mais apertados para cenários Fora da Distribuição (OOD). A restrição conservativa atua como um regularizador geométrico, impedindo que o modelo aprenda direções vetoriais arbitrárias e forçando-o a extrapolar a forma da paisagem de energia, levando a uma melhor robustez sob deslocamentos de distribuição.
4. Identificabilidade e Propagação de Erro: O artigo estabelece que, embora o valor Q absoluto não possa ser unicamente identificado (devido à constante dependente do estado desconhecida), o ranking de ações dentro do estado é exato. Além disso, limita a propagação de erros de estimativa de pontuação para preferências de ação, mostrando que os erros escalam linearmente com a distância entre as ações (continuidade de Lipschitz).

4. Resultados Experimentais

Os autores avaliaram o ENERGYFLOW nos benchmarks RoboMimic e Meta-World, bem como em tarefas de robôs do mundo real.

• Desempenho de Imitação (RQ1): O ENERGYFLOW alcançou desempenho State-of-the-Art (SOTA).
  • RoboMimic: 93,8% de taxa de sucesso média (vs. 91,2% para Diffusion Policy).
  • Meta-World: 92,5% de taxa de sucesso média (vs. 90,7% para Diffusion Policy).
  • Superou significativamente outros métodos baseados em energia (como Implicit BC) e baselines de IRL.
• Implantação no Mundo Real (RQ2): Implantado em um robô AGIBOT G1 para tarefas "Garrafa" e "Gaveta". O ENERGYFLOW alcançou 100% de taxa de sucesso com trajetórias suaves e menos hesitações em comparação com as baselines.
• Qualidade da Recompensa para RL a Jusante (RQ3): Quando usado como sinal de recompensa para treinar um agente Soft Actor-Critic (SAC):
  • A recompensa "Energia Centralizada" permitiu que o SAC alcançasse desempenho próximo ao oráculo, superando significativamente recompensas esparsas e recompensas de energia bruta.
  • Forneceu orientação densa, passo a passo, que preveniu platôs iniciais no treinamento.
• Generalização OOD (RQ4): Sob perturbações crescentes na posição inicial, o ENERGYFLOW manteve taxas de sucesso robustas onde políticas de difusão e fluxo não restritas degradaram-se significativamente. Isso confirma o benefício teórico da restrição conservativa.
• Eficiência (RQ6): O ENERGYFLOW alcançou velocidades de inferência comparáveis às Políticas de Fluxo não baseadas em energia (~10ms de latência) enquanto fornecia os benefícios de uma função de energia explícita, evitando os tempos de amostragem lentos de Modelos Baseados em Energia (EBMs) tradicionais.

5. Significado e Impacto

• Unificação de Imitação e RL: O ENERGYFLOW preenche a lacuna entre Clonagem de Comportamento e Aprendizado por Reforço. Permite que uma política seja treinada via imitação enquanto expõe simultaneamente um sinal de recompensa válido para refinamento posterior, eliminando a necessidade de design manual de recompensa ou IRL adversarial instável.
• Garantias Teóricas para Difusão: O trabalho fornece a primeira justificativa teórica para por que impor um campo conservativo (integrabilidade) em políticas de difusão é benéfico, ligando-o à redução da complexidade de hipóteses e melhoria na generalização OOD.
• Utilidade Prática: O método é computacionalmente eficiente, escalável para espaços de ação de alta dimensão e eficaz em configurações robóticas do mundo real, oferecendo uma solução robusta para aprender habilidades complexas de manipulação a partir de demonstrações limitadas.

Em resumo, o ENERGYFLOW demonstra que as restrições estruturais necessárias para extração válida de recompensa (campos conservativos) servem simultaneamente como vieses indutivos poderosos que aprimoram a generalização e a robustez de políticas baseadas em difusão.