Resumo Técnico: Correspondência Q-Adjoint com Região de Confiança (TRQAM)

Declaração do Problema

O artigo aborda o desafio do ajuste fino off-policy para políticas de fluxo pré-treinadas. Embora as políticas de correspondência de fluxo ofereçam distribuições de ação ricas e multi-modais superiores às políticas Gaussianas unimodais convencionais, sua definição implícita por meio de processos de remoção de ruído em múltiplos passos torna a melhoria de política baseada em gradiente instável. A retropropagação direta através da cadeia de amostragem é computacionalmente cara e propensa a instabilidade.

Soluções existentes tentam contornar isso aprendendo resíduos aditivos (ignorando a dinâmica de múltiplos passos) ou operando no espaço de ruído (limitado pela expressividade da política congelada). Uma abordagem recente, Aprendizado por Reforço Q com Correspondência Adjoint (QAM), reformula o ajuste fino como um problema de Controle Ótimo Estocástico (SOC) sem memória, utilizando um crítico aprendido para guiar o processo de amostragem via correspondência adjoint. No entanto, os autores identificam uma fragilidade fundamental no QAM: em configurações off-policy, o crítico aprendido é inevitavelmente imperfeito. Pequenos erros de aproximação no crítico são amplificados exponencialmente quando a atualização da política depende de uma "temperatura" fixa (ou temperatura inversa $\beta$), levando a um desvio destrutivo do prior pré-treinado e ao colapso eventual do modelo. Essa instabilidade persiste mesmo com clipping de gradiente.

Metodologia: Correspondência Q-Adjoint com Região de Confiança (TRQAM)

O TRQAM propõe um algoritmo de ajuste fino off-policy estável que integra um mecanismo de região de confiança diretamente nas dinâmicas de amostragem do SOC. A inovação central é o controle adaptativo da divergência Kullback-Leibler (KL) no espaço de caminhos entre a política ajustada e o prior pré-treinado.

Principais Insights Teóricos

1. Amplificação Exponencial de Erros: Os autores formalizam (Lema 1) que a correspondência adjoint de temperatura fixa sofre de uma amplificação exponencial de erros do crítico. Um $\beta$ fixo não pode simultaneamente explorar um crítico confiável e proteger contra um não confiável.
2. $\lambda$ como Parâmetro de Região de Confiança: O TRQAM introduz um parâmetro escalar $\lambda$ que escala o coeficiente de difusão na EDE do SOC por $\sqrt{\lambda}$.
3. Identidade KL Exata (Teorema 1): Usando o teorema de Girsanov, os autores provam que escalar a difusão por $\sqrt{\lambda}$ torna a divergência KL no espaço de caminhos uma função exata e de forma fechada de $\lambda$:
   $$D_{KL}(P_u \parallel P_{base}) = \mathbb{E}_{X \sim P_u} \left[ \frac{1}{2\lambda} \int_0^1 \|u(X_\tau, \tau)\|^2 d\tau \right]$$
   Isso estabelece um vínculo estrutural direto entre o parâmetro $\lambda$ e o desvio da política base.
4. Limite Terminal KL (Proposição 1): A KL no espaço de caminhos limita superiormente a KL da distribuição de ação terminal, garantindo que controlar a KL no espaço de caminhos restrinja efetivamente o desvio final da política.

Estrutura Algorítmica

Em vez de adicionar um termo KL como uma penalidade de perda externa (que pode ser anulada por gradientes fortes do crítico), o TRQAM internaliza $\lambda$ dentro das dinâmicas de amostragem:

• Amostragem: A EDE controlada usa um coeficiente de difusão $\sqrt{\lambda}\sigma(\tau)$. Ajustar $\lambda$ remodela a própria EDE, puxando as dinâmicas controladas mais perto ou mais longe das dinâmicas base.
• Adaptação: O algoritmo emprega descida dual projetada para adaptar $\lambda$ dinamicamente. Ele estima a KL no espaço de caminhos realizada ($\hat{D}_n$) via um substituto de Monte Carlo derivado do amostrador sem memória OT discretizado.
• Regra de Atualização: $\lambda$ é atualizado para impor um orçamento de KL alvo $\epsilon_{KL}$:
  $$\lambda_{n+1} \leftarrow \max\{0, \lambda_n + \eta_\lambda (D_n - \epsilon_{KL})\}$$
  Se a KL realizada exceder o orçamento, $\lambda$ aumenta, tornando as dinâmicas mais conservadoras; se cair abaixo, $\lambda$ diminui, permitindo uma melhoria mais agressiva.

Principais Contribuições

1. Identificação de Fragilidade: O artigo identifica e formaliza a amplificação exponencial de erros do crítico na correspondência adjoint de temperatura fixa, demonstrando sua tendência a causar colapso do modelo em benchmarks como Robomimic.
2. Vínculo Teórico: Prova que escalar o coeficiente de difusão por $\sqrt{\lambda}$ cria uma identidade exata entre o custo de controle e a KL no espaço de caminhos, transformando $\lambda$ em um parâmetro de região de confiança principiado.
3. Design do Algoritmo: O TRQAM é proposto como um método que internaliza o parâmetro de região de confiança nas dinâmicas do SOC e o adapta via descida dual, impondo o limite KL no nível de amostragem em vez de como uma penalidade de perda suave.
4. Superioridade Empírica: Experimentos extensivos demonstram que o TRQAM alcança convergência estável onde métodos anteriores falham e supera as linhas de base mais avançadas.

Resultados Experimentais

Os autores avaliam o TRQAM em 50 tarefas do benchmark OGBench e no benchmark de manipulação Robomimic.

• Desempenho de RL Offline: No OGBench, o TRQAM alcança uma taxa de sucesso geral de 68%, superando substancialmente a linha de base mais forte (DSRL em 46%) e os métodos anteriores de correspondência adjoint (QAM em 35%, QAM-E em 45%). Os ganhos são particularmente pronunciados em tarefas de horizonte longo e combinatórias.
• Estabilidade: No Robomimic, métodos de temperatura fixa (QAM, QAM-E) exibem perdas adjoint explodindo (até $10^{20}$) e taxas de sucesso colapsando para próximo de zero em várias configurações de hiperparâmetros. Em contraste, o TRQAM permanece estável em todos os orçamentos testados.
• Validação do Mecanismo: Estudos de ablação confirmam que:
  • A adaptação é necessária: Orçamentos KL adaptativos superam os fixos.
  • A internalização é necessária: Internalizar $\lambda$ na EDE (TRQAM) rastreia estritamente o orçamento KL, enquanto a regularização KL externa permite que a KL realizada desvie muito além do alvo, levando à degradação de desempenho.
• Sensibilidade: O orçamento de KL alvo $\epsilon_{KL}$ atua como um controle previsível; orçamentos mais apertados geralmente produzem melhor desempenho, embora o valor ótimo varie sistematicamente com a estrutura da tarefa (por exemplo, orçamentos maiores são benéficos para tarefas com espaços de estado maiores, como puzzle-4x4).

Significado e Alegações

O artigo afirma que o TRQAM fornece uma alternativa principiada e estável aos métodos existentes de ajuste fino off-policy para políticas de fluxo. Ao internalizar a restrição de região de confiança nas dinâmicas de amostragem via teorema de Girsanov, o TRQAM evita o "desvio destrutivo" causado por críticos imperfeitos em métodos como o QAM.

Os autores posicionam o TRQAM como o análogo off-policy da Otimização de Política com Região de Confiança (TRPO) e da Otimização de Política Próxima (PPO) para RL on-policy. Eles argumentam que, para políticas de fluxo pré-treinadas, que servem como priors comportamentais valiosos, manter um desvio controlado via uma identidade KL exata é crucial para prevenir o esquecimento catastrófico e o colapso do modelo. O método é apresentado como uma solução robusta tanto para RL offline quanto para a transição para ajuste fino online, oferecendo um mecanismo confiável para aproveitar habilidades pré-treinadas sem sacrificar a estabilidade.