DAWM: Diffusion Action World Models for Offline Reinforcement Learning via Action-Inferred Transitions

Resumo Técnico: DAWM – Modelos de Mundo de Ação por Difusão para Aprendizado por Reforço Offline via Transições Inferidas por Ação

1. Declaração do Problema

Modelos de mundo baseados em difusão demonstraram capacidades robustas na síntese de trajetórias realistas de longo horizonte para aprendizado por reforço (RL) offline. No entanto, os métodos existentes enfrentam um problema crítico de compatibilidade com algoritmos padrão de RL offline baseados em valor (por exemplo, TD3BC, IQL) que dependem de aprendizado de diferença temporal (TD) de um passo.

• Limitação de Modelos Livres de Ação: Métodos como o Modelo de Mundo por Difusão (DWM) geram estados e recompensas futuras, mas não geram ações explicitamente. Consequentemente, eles não podem construir transições completas $(s, a, r, s')$, tornando-os incompatíveis com o aprendizado TD de um passo, que exige consistência entre estados, ações e recompensas para uma regularização eficaz do crítico.
• Limitação da Modelagem Conjunta: Abordagens que tentam modelar conjuntamente estados, recompensas e ações frequentemente sofrem com aumento da complexidade de treinamento e instabilidade, levando a uma redução no desempenho prático.
• Limitação de Métodos Baseados em Planejamento: Planejadores baseados em difusão (por exemplo, Diffuser, Decision Diffuser) que recuperam ações por meio de Modelos de Dinâmica Inversa (IDM) ou modelagem conjunta geralmente dependem de amostragem iterativa e seleção de recompensa post-hoc. Isso resulta em altos custos de inferência e contorna a eficiência do aprendizado TD de um passo.

O desafio central é gerar dados sintéticos que suportem aprendizado TD de um passo estável e eficiente, sem incorrer na instabilidade de treinamento da modelagem conjunta ou nos custos de inferência do planejamento iterativo.

2. Metodologia: DAWM

Os autores propõem o DAWM (Modelo de Mundo de Ação por Difusão), um framework modular que desacopla a geração de trajetórias da inferência de ação para produzir transições completas e rotuladas adequadas para RL offline. O framework consiste em duas fases principais:

Fase 1: Geração Condicional de Trajetórias por Difusão

• Um modelo de difusão $p_\theta$ é treinado para gerar sequências futuras de estados e recompensas condicionadas ao estado atual $s_t$, ação $a_t$ e um retorno-alvo ($R_{tg}$).
• O modelo define a distribuição condicional $p_\theta(\hat{s}_{t+1:t+H+1}, \hat{r}_{t:t+H} \mid s_t, a_t, R_{tg})$.
• O objetivo de treinamento minimiza o erro esperado entre o ruído previsto e o ruído verdadeiro usando um objetivo de remoção de ruído com pontuação e orientação livre de classificador.
• Crucialmente, este modelo gera trajetórias livres de ação (apenas estados e recompensas), simplificando a tarefa generativa ao evitar as complexidades de modelar ações conjuntamente dentro do processo de difusão.

Fase 2: Completamento de Ação via Modelo de Dinâmica Inversa (IDM)

• Para converter as trajetórias livres de ação em transições completas, um Modelo de Dinâmica Inversa (IDM) separado, $f_\phi$, treinado independentemente, é empregado.
• O IDM reconstrói as ações faltantes modelando a distribuição $f_\phi(a_t \mid s_{t-m:t-1}, s_t, s_{t+1})$ como uma distribuição Gaussiana diagonal.
• O IDM é treinado em dados offline reais para maximizar a verossimilhança das ações observadas dadas as transições de estado.
• Na inferência, o IDM é aplicado às trajetórias geradas pela difusão para inferir ações, produzindo tuplas completas $(s, a, r, s')$.

Treinamento de Política Offline

• As transições completas sintetizadas são usadas para treinar agentes de RL offline usando algoritmos padrão baseados em valor, especificamente TD3BC (que adiciona um regularizador de clonagem de comportamento ao TD3) e IQL (que realiza regressão supervisionada ponderada por vantagem).
• As funções Q são atualizadas usando alvos de Bellman computados a partir das transições reconstruídas, permitindo atualizações TD de um passo eficientes.

3. Contribuições Principais

1. Framework Modular: Os autores propõem um framework eficiente que combina geração de estado-recompensa baseada em difusão com um IDM treinado separadamente para inferência de ação. Este design evita a instabilidade de treinamento da modelagem conjunta enquanto produz transições completas.
2. Compatibilidade com Aprendizado TD de Um Passo: Ao gerar transições completas $(s, a, r, s')$, o DAWM permite o uso de algoritmos conservadores de RL offline (TD3BC, IQL) que dependem de aprendizado TD de um passo, oferecendo uma alternativa computacionalmente eficiente aos planejadores de difusão iterativos.
3. Desempenho Empírico: O artigo demonstra que algoritmos conservadores de RL offline se beneficiam significativamente do treinamento nessas trajetórias aumentadas, superando consistentemente bases anteriores baseadas em difusão (DWM, DD) em múltiplas tarefas de locomoção D4RL.
4. Viabilidade de Dados Sintéticos: Os resultados indicam que agentes treinados em dados sintéticos gerados pelo DAWM podem alcançar desempenho comparável ou superior ao de agentes treinados em conjuntos de dados offline reais, destacando o potencial de trajetórias sintéticas como uma alternativa prática.

4. Resultados Experimentais

Os experimentos foram conduzidos em 9 benchmarks padrão de locomoção D4RL (Hopper, Walker2d, HalfCheetah) com qualidade de dados variável (Medium, Medium-Replay, Medium-Expert).

• Comparação com DWM: O DAWM superou consistentemente a base DWM livre de ação. Com TD3BC, o DAWM alcançou uma melhoria média de 9,3%, e com IQL, 9,5%. Os ganhos foram particularmente pronunciados em tarefas de controle complexas como Walker2d.
• Comparação com Métodos Baseados em Planejamento (DD): O DAWM superou o Decision Diffuser (DD) em termos de retorno normalizado. Além disso, o DAWM demonstrou eficiência significativamente maior, alcançando uma inferência aproximadamente 4,0 vezes mais rápida em comparação com o DD, pois evita a amostragem iterativa durante a inferência.
• Comparação com Dados Reais: Agentes treinados em dados gerados pelo DAWM alcançaram retornos normalizados comparáveis aos de agentes treinados em dados reais usando TD3BC/IQL padrão e, em alguns casos (por exemplo, Hopper-MR com TD3BC), superaram-nos. Embora o SSORL (uma base que usa IDM em dados reais) tenha se saído bem, os autores observam que o SSORL se beneficia da filtragem para as trajetórias mais longas, um fator não presente na geração de comprimento fixo do DAWM.
• Estudos de Ablação:
  • Volume de Dados (DAWM-T): Reduzir o número de pares de condicionamento por um fator de 8 (DAWM-T) resultou em quedas de desempenho apenas modestas, sugerindo que o mecanismo de completamento de ação é um fator mais crítico do que o volume bruto de dados.
  • Horizonte de Geração: O DAWM manteve desempenho robusto em diferentes horizontes de geração ($H=1, 3, 7$), indicando que o framework de aprendizado TD de um passo é estável independentemente do comprimento da simulação.

5. Significado e Alegações

O artigo afirma que o DAWM preenche a lacuna entre modelagem generativa eficiente e aprendizado de política baseado em TD. Seu significado principal reside em:

• Habilitar Aprendizado TD de Um Passo: Fornece um mecanismo para usar dados gerados por difusão em algoritmos padrão e estáveis de RL offline que anteriormente eram incompatíveis com modelos de mundo livres de ação.
• Eficiência Computacional: Oferece uma solução prática para RL generativo offline em grande escala ao reduzir os custos de inferência em comparação com métodos de planejamento iterativo.
• Alternativa Prática: Os resultados sugerem que trajetórias sintéticas geradas por essa abordagem modular podem servir como uma alternativa viável e de alta qualidade a conjuntos de dados offline reais, potencialmente mitigando problemas de escassez de dados em aplicações do mundo real onde a exploração é cara ou insegura.

Os autores concluem que, embora o método mostre promessa, trabalhos futuros investigarão sua extensão para controle visual, robótica de alta dimensão e robustez sob dinâmicas imperfeitas.