Resumo Técnico: Direcionamento de Difusão por Perturbação Lagrangiana (LP-DS)

1. Definição do Problema

Políticas generativas de alta capacidade, como as baseadas em modelos de difusão e flow-matching, estabeleceram-se como paradigmas poderosos para controle contínuo e robótica. Esses modelos são excelentes na representação de distribuições de ações expressivas e multimodais aprendidas a partir de dados offline. No entanto, eles enfrentam duas limitações primárias:

1. Restrições de Clonagem de Comportamento (BC): Políticas treinadas via BC herdam as limitações de seus conjuntos de dados de demonstração, incluindo cobertura incompleta do espaço de estados, a omissão de comportamentos de recuperação raros e fragilidade sob mudança de distribuição.
2. Instabilidade de Aprendizado por Reforço (RL): Embora o RL ofereça um mecanismo para melhoria online, o ajuste fino direto de decodificadores generativos de grande escala é frequentemente instável e ineficiente em termos de amostras. A interação entre os objetivos de otimização e a dinâmica de denoising ou integração de longo horizonte frequentemente leva à instabilidade de treinamento.

Alternativas existentes, como o Direcionamento de Difusão via Aprendizado por Reforço (DSRL), tentam melhorar políticas congeladas aprendendo uma política de ruído latente. No entanto, o DSRL efetivamente substitui o prior de pré-treinamento por uma política latente aprendida. Essa abordagem sofre de dois modos de falha:

• Deriva Fora da Variedade (Off-Manifold Drift): O ruído direcionado pode derivar para longe do suporte Gaussiano padrão usado para treinar o decodificador, levando a uma decodificação instável.
• Colapso de Modo (Mode Collapse): A política latente aprendida pode colapsar o comportamento multimodal do backbone para um subconjunto estreito de modos, reduzindo a diversidade comportamental.

2. Metodologia: Direcionamento de Difusão por Perturbação Lagrangiana (LP-DS)

Os autores propõem o LP-DS, um framework de adaptação leve que melhora políticas generativas congeladas através do aprendizado de um resíduo condicionado ao estado no espaço de ruído latente, em vez de modificar os pesos do decodificador ou substituir inteiramente o prior.

Mecanismo Central

O LP-DS trata o backbone generativo congelado $\Phi$ como um decodificador de caixa-preta que mapeia o ruído latente $w$ e o estado $s$ para ações. Em vez de amostrar o ruído diretamente de um prior $p_0$ (ex: $\mathcal{N}(0, I)$), o LP-DS introduz uma rede de perturbação aprendível $\Delta_\theta(s)$:
$$w = \epsilon + \Delta_\theta(s), \quad \text{onde } \epsilon \sim \mathcal{N}(0, I)$$
Quando $\Delta_\theta(s) = 0$, a política recupera exatamente a distribuição original de clonagem de comportamento.

Otimização Restrita via Relaxação Lagrangiana

Para evitar que a perturbação empurre o ruído para fora da variedade onde o decodificador foi treinado, o LP-DS formula o problema de aprendizado como uma otimização restrita. O objetivo é maximizar o valor Q esperado das ações geradas sujeito a uma restrição de região de confiança sobre a magnitude da perturbação.

Os autores aproximam a divergência Kullback–Leibler (KL) esperada entre a distribuição de ruído direcionada e o prior usando a magnitude ao quadrado do deslocamento:
$$\mathbb{E}_{s \sim D}[\|\Delta_\theta(s)\|^2_2] \leq \delta$$
onde $\delta$ é um hiperparâmetro que controla o desvio permitido.

Isso é resolvido usando uma abordagem de relaxação Lagrangiana. Uma variável dual (multiplicador de Lagrange) $\alpha$ é introduzida para impor a restrição dinamicamente:
$$\mathcal{L}(\theta, \alpha) = \mathbb{E}[Q^W(s, w) - \alpha(\|\Delta_\theta(s)\|^2_2 - \delta)]$$
O algoritmo alterna entre:

1. Atualização do Ator: Maximizando $\mathcal{L}$ em relação a $\theta$ (ascensão de gradiente).
2. Atualização Dual: Maximizando $\mathcal{L}$ em relação a $\alpha$ (ascensão de gradiente projetada) para garantir $\alpha \geq 0$. Se a perturbação violar a região de confiança, $\alpha$ aumenta, penalizando o ator e forçando-o a aderir mais proximamente ao prior.

Componentes do Algoritmo

O procedimento de treinamento (Algoritmo 1) envolve:

• Coleta de Dados: Amostragem de ações usando a política direcionada $w = \epsilon + \Delta_\theta(s)$.
• Atualizações do Crítico: Atualização de um crítico de valor de ação $Q^A$ via aprendizado de Diferença Temporal (TD) e um crítico de valor latente $Q^W$ via destilação de $Q^A$ na distribuição de ruído base.
• Atualizações Primal-Dual: Atualização da rede de perturbação $\theta$ e do multiplicador $\alpha$ para equilibrar a maximização de recompensa com a preservação do prior.

3. Principais Contribuições

O artigo descreve quatro contribuições principais:

1. Framework LP-DS: Um método de adaptação leve que melhora políticas generativas congeladas através do controle de ruído residual condicionado ao estado, evitando a instabilidade e o custo do ajuste fino total do decodificador.
2. Objetivo de Região de Confiança Lagrangiana: Um objetivo inovador que restringe dinamicamente a magnitude da perturbação, fornecendo um mecanismo controlável para equilibrar a maximização da recompensa contra a preservação do prior latente pré-treinado.
3. Mitigação de Colapso de Modo: Evidência empírica mostrando que o LP-DS preserva maior diversidade no espaço de ações e mitiga o colapso de modo em comparação com métodos de direcionamento no espaço de ruído que substituem o prior (como o DSRL).
4. Generalização: Validação do método além de políticas de difusão compactas, incluindo backbones de flow-matching, um modelo de visão-linguagem-ação de grande escala ($\pi_0$) e experimentos físicos com o robô Franka.

4. Resultados Experimentais

Os autores avaliaram o LP-DS nos benchmarks RoboMimic (manipulação), OpenAI Gym (locomoção) e Adroit (manipulação destreza), comparando-o com DSRL, DPPO (ajuste fino de gradiente de política) e baselines de Q-learning de difusão offline-to-online.

• Desempenho: O LP-DS melhorou consistentemente as taxas de sucesso e os retornos episódicos. Notavelmente, no ambiente WALKER2D, alcançou aproximadamente um melhoria de 25% no retorno sobre o baseline mais forte.
• Eficiência de Amostra: Ao atualizar apenas um módulo de perturbação leve em vez do decodificador completo, o LP-DS demonstrou superior eficiência de amostra e dinâmica de aprendizado mais estável.
• Preservação de Diversidade:
  • Em uma tarefa de navegação de múltiplos objetivos simplificada, o LP-DS permitiu o controle explícito do trade-off entre desempenho e diversidade via o limite da região de confiança $\delta$. Valores menores de $\delta$ preservaram o comportamento multimodal, enquanto valores maiores permitiram especialização. Em contraste, o DSRL colapsou imediatamente para um único modo.
  • No ambiente AVOIDING, o LP-DS preservou um amplo conjunto de rotas viáveis de desvio de obstáculos (multimodalidade de nível de trajetória) com $\delta$ pequeno, enquanto o DSRL colapsou para um único caminho dominante.
  • Análise de Entropia: O LP-DS manteve consistentemente uma entropia de espaço de ação mais alta que o DSRL durante a adaptação online, indicando que a restrição Lagrangiana desencoraja efetivamente o colapso prematuro enquanto ainda melhora o desempenho.
• Estudos de Ablação: Remover a penalidade Lagrangiana levou a um aprendizado instável e degradou o sucesso. Remover tanto a penalidade Lagrangiana quanto o limite de ruído resultou em comportamento altamente instável e baixas taxas de sucesso, confirmando que a restrição de região de confiança é o principal mecanismo para prevenir o direcionamento fora da distribuição.
• Sensibilidade: O método mostrou-se robusto ao valor específico do alvo de região de confiança $\delta$, com valores acima de 0,1 geralmente apresentando forte desempenho.

5. Significância e Alegações

O artigo afirma que o LP-DS oferece uma alternativa estável e eficiente para o ajuste fino total de políticas generativas. Ao tratar o backbone congelado como um prior fixo e aprender apenas um deslocamento residual limitado, o método:

• Previne a Deriva Fora da Variedade: Mantém as consultas de ruído dentro do intervalo operacional válido do decodificador, evitando o comportamento errático visto no direcionamento de ruído não restrito.
• Preserva a Multimodalidade: Mitiga o colapso de modo frequentemente observado ao substituir o prior de pré-treinamento, mantendo assim a diversidade comportamental inerente ao backbone congelado.
• Generaliza entre Arquiteturas: A abordagem não se limita a modelos de difusão, mas aplica-se a qualquer política generativa com um mapeamento latente-para-ação, incluindo modelos de flow-matching e modelos de visão-linguagem-ação de grande escala.
• Permite Implantação Física: O método foi validado com sucesso em experimentos físicos com o robô Franka usando adaptação latente baseada em simulação, demonstrando sua aplicabilidade além de benchmarks simulados.

Os autores posicionam o LP-DS como uma solução prática para a "melhoria online" de políticas generativas de alta capacidade, equilibrando a necessidade de otimização específica da tarefa com a necessidade de manter a estrutura multimodal robusta do modelo pré-treinado.