EXPO: Stable Reinforcement Learning with Expressive Policies

1. Declaração do Problema

O artigo aborda o desafio de ajustar finamente classes de políticas expressivas (como Políticas de Difusão e Políticas de Correspondência de Fluxo) usando Aprendizado por Reforço (RL) Online dado um conjunto de dados offline.

• O Conflito Central: Embora políticas expressivas se destaquem em capturar distribuições de comportamento complexas e multimodais via Aprendizado por Imitação (IL), elas são difíceis de otimizar diretamente com RL. Diferentemente de políticas Gaussianas simples, políticas expressivas são parametrizadas por longas cadeias de remoção de ruído. A retropropagação de gradientes de valor através dessas cadeias para atualizar os parâmetros da política é computacionalmente cara e instável, frequentemente levando ao colapso do treinamento ou à convergência pobre.
• A Lacuna: Métodos existentes de RL online são projetados para políticas Gaussianas simples e falham em aproveitar a capacidade de modelos expressivos pré-treinados. Por outro lado, métodos que tentam ajustar finamente políticas expressivas frequentemente lutam com estabilidade ou exigem modificações complexas no processo de remoção de ruído.
• Objetivo: Desenvolver um algoritmo de RL online eficiente em amostras e estável que possa ajustar finamente políticas expressivas pré-treinadas usando dados offline sem otimizar diretamente os parâmetros da política expressiva contra a função de valor.

2. Metodologia: EXPO (Otimização de Política Expressiva)

Os autores propõem o EXPO, um framework que desacopla a política expressiva da maximização direta de valor. Em vez de retropropagar gradientes através da política expressiva, o EXPO constrói uma política "em tempo real" (OTF) que maximiza valores-Q usando dois componentes distintos:

A. Arquitetura de Política de Dois Componentes

1. Política Expressiva Base ($\pi_{base}$):
   • Função: Gera candidatos de ação iniciais baseados na distribuição de comportamento.
   • Treinamento: Treinada usando um objetivo de Aprendizado por Imitação (IL) estável (por exemplo, Modelos Probabilísticos de Difusão com Remoção de Ruído). Ela nunca é explicitamente treinada para maximizar o valor-Q, preservando sua estabilidade e capacidade de capturar modos complexos.
2. Política de Edição Leve ($\pi_{edit}$):
   • Função: Uma pequena política Gaussiana que recebe uma ação $a$ amostrada de $\pi_{base}$ e produz uma edição residual $\hat{a}$ para gerar uma ação refinada $\tilde{a} = a + \hat{a}$.
   • Treinamento: Treinada usando gradientes de política de RL padrão (maximizando $Q(s, a + \hat{a})$) com regularização de entropia.
   • Restrição: A edição é restringida (escalada por um hiperparâmetro $\beta$) para permanecer próxima da ação original, garantindo que a política de edição resolva um problema de otimização local em vez de global. Isso permite um treinamento eficiente e estável.

B. Extração de Política em Tempo Real (OTF)

A política real usada para amostragem e backup de Diferença Temporal (TD) não é uma única rede neural, mas um processo de seleção dinâmico:

1. Amostragem: Para um estado dado $s$, o sistema amostra $N$ ações de $\pi_{base}$.
2. Edição: Cada ação amostrada é refinada por $\pi_{edit}$.
3. Seleção: O sistema avalia o valor-Q tanto para as ações originais ($a_i$) quanto para as ações editadas ($\tilde{a}_i$).
4. Maximização: A ação com o maior valor-Q entre todos os candidatos é selecionada para execução e usada como alvo para o backup TD.

C. Backup de Entropia para Regimes com Dados Limitados

Em cenários com pequenos conjuntos de dados offline, os autores introduzem um backup de entropia. Como políticas expressivas (como difusão) carecem de entropia de forma fechada, eles constroem uma distribuição de amostragem suave sobre os $N$ candidatos de ação base+editados, ponderados por seus valores-Q. Isso permite a inclusão de um bônus de entropia na perda de RL para encorajar a exploração sem exigir cálculos analíticos de entropia.

3. Contribuições Principais

1. Estratégia de Maximização de Valor Estável: O artigo introduz uma abordagem inovadora para evitar a propagação direta de gradientes através de cadeias de remoção de ruído expressivas. Ao transferir a maximização de valor para uma política de edição leve e um mecanismo de seleção, o EXPO alcança estabilidade onde a otimização direta falha.
2. Eficiência em Amostras: O EXPO alcança uma melhoria de 2 a 3 vezes na eficiência de amostras em comparação com métodos anteriores do estado da arte em configurações de RL online e de ajuste fino offline-para-online.
3. Agnosticismo de Política: O framework é agnóstico à parametrização específica da política base (funciona com Difusão, Correspondência de Fluxo, etc.) e não requer uma função de valor pré-treinada, tornando-o aplicável a qualquer política robótica pré-treinada.
4. Backup TD em Tempo Real: Os autores demonstram que o uso da ação que maximiza o valor para o alvo TD (não apenas para amostragem) é crítico para a velocidade de aprendizado, efetivamente fechando a lacuna entre atualizações semelhantes a SARSA e semelhantes a Q-learning no contexto de políticas expressivas.

4. Resultados Experimentais

O método foi avaliado em 12 tarefas desafiadoras de controle contínuo em 4 domínios (AntMaze, Adroit, Robomimic, MimicGen), todos com recompensas esparsas.

• RL Online (Sem Pré-treinamento): O EXPO superou significativamente as linhas de base (incluindo RLPD, IDQL e QSM) na eficiência de amostras, frequentemente convergindo mais rápido mesmo sem aproveitar o conjunto de dados offline para pré-treinamento.
• Ajuste Fino Offline-para-Online:
  • O EXPO ajustou com sucesso políticas pré-treinadas sem o colapso de desempenho frequentemente observado em outros métodos (por exemplo, o DAC colapsou, o IDQL estagnou).
  • Mantive alto desempenho em tarefas de manipulação de alta granularidade (por exemplo, enfiar uma agulha, empilhar blocos) onde outros métodos lutaram para melhorar além do desempenho inicial de aprendizado por imitação.
• Estudos de Ablação:
  • Backup TD OTF: Remover a seleção que maximiza o valor para o alvo TD (usando apenas amostragem) reduziu drasticamente o desempenho, confirmando a importância do alinhamento de valor-Q no backup.
  • Edições de Ação: Remover a política de edição resultou em desempenho pobre, especialmente em tarefas que exigem exploração além da estreita distribuição de comportamento.
  • Qualidade dos Dados Offline: O desempenho correlacionou-se fortemente com a qualidade do conjunto de dados offline. No entanto, a variante de backup de entropia permitiu que o EXPO tivesse sucesso mesmo com dados offline muito ruins (onde o aprendizado por imitação sozinho alcançou <10% de sucesso).

5. Significado

O EXPO representa um avanço significativo em Robótica e RL ao resolver o gargalo do "ajuste fino de política expressiva".

• Impacto Prático: Permite a implantação de grandes modelos generativos pré-treinados (que são excelentes em imitar dados humanos) em ambientes do mundo real onde devem se adaptar e melhorar via interação online.
• Insight Teórico: Desafia a necessidade de retropropagar através de cadeias generativas complexas para RL, mostrando que uma estratégia de "edição local + seleção global" é mais estável e eficiente.
• Direção Futura: O trabalho abre caminho para o ajuste fino de qualquer política generativa pré-treinada (difusão, fluxo, etc.) em configurações online, potencialmente acelerando a adoção de modelos de fundação em robótica.

Limitações: O método requer a amostragem de múltiplas ações para cada atualização de lote, o que é computacionalmente caro. Além disso, assume-se que existe um prior razoável (conjunto de dados offline ou política pré-treinada); não é projetado para aprendizado do zero sem conhecimento prévio.