Coupled Local and Global World Models for Efficient First Order RL

Resumo Técnico: Modelos de Mundo Acoplados Local e Global para RL de Primeira Ordem Eficiente

Declaração do Problema

O Aprendizado por Reforço (RL) alcançou um sucesso significativo em tarefas de locomoção via simuladores paralelos em GPU, mas enfrenta barreiras substanciais no controle perceptual e na manipulação. Dois problemas primários dificultam o progresso:

1. Ineficiência de Amostragem: O aprendizado online prático diretamente de interações com robôs reais é frequentemente inviável devido ao alto custo da coleta de dados.
2. Fidelidade de Simulação: Simular tarefas de manipulação é inerentemente mais complexo do que a locomoção. O estado do ambiente muda dinamicamente com as ações do agente (ex: objetos deformáveis, dinâmica granular), introduzando uma variabilidade infinita que os simuladores de física atuais lutam para modelar com fidelidade suficiente para a transferência sim-to-real.

Embora o aprendimento de agentes dentro de modelos de mundo diferenciáveis usando métodos de Gradiente de Primeira Ordem (FoG) ofereça um caminho para resolver a ineficiência de amostragem, as abordagens existentes enfrentam um gargalo de escalabilidade. Modelos de mundo de grande escala (ex: modelos de difusão) necessários para dinâmica de mundo real de alta fidelidade são computacionalmente caros para serem diferenciados, tornando a retropropagação padrão intratável para o treinamento de políticas FoG. Por outro via, modelos leves adequados para diferenciação carecem de expressividade para capturar dinâmicas complexas do mundo real.

Metodologia

O artigo propõe um novo framework que desacopla a simulação direta da diferenciação reversa para permitir o RL baseado em modelo de Primeira Ordem (FoG-MBRL) eficiente usando modelos de mundo de grande escala. A arquitetura central consiste em dois componentes acoplados:

1. Modelo de Mundo Global (Passagem Direta/Forward Pass):

   • Um modelo de grande escala, treinado offline (seja um modelo baseado em difusão como o DIAMOND ou um modelo de difusão latente baseado em transformer como o DreamerV4), atua como um simulador "global".
   • Este modelo opera no espaço de pixels (ou espaço latente de alta dimensão) para gerar trajetórias diretas de alta fidelidade ($s^{\rightarrow}_{t+1}$) que capturam com precisão complexas dinâmicas visuais e de contato.
   • Ele é treinado em dados de exploração (play data) não estruturados e de alta entropia coletados de robôs reais, eliminando a necessidade de simuladores de física feitos à mão.

2. Modelo de Mundo Local (Passagem Reversa/Backward Pass):

   • Um Modelo de Espaço de Estado Recorrente (RSSM) leve atua como um substituto (surrogate) "local" em um espaço latente de baixa dimensão.
   • Este modelo aproxima as dinâmicas locais e fornece os gradientes necessários para a otimização da política.
   • Crucialmente, os gradientes são computados avaliando os Jacobianos do modelo local nos pontos de estado gerados pelo modelo global ($s^{\rightarrow}_{t+1}$), em vez das próprias previsões do modelo local. Esse desacoplamento permite que o sistema aproveite a precisão do modelo global para a geração de trajetórias enquanto utiliza o modelo local para uma estimativa de gradiente computacionalmente tratável.

3. Protocolo de Treinamento:

   • Coleta de Dados: Dados de exploração não estruturados são coletados em robôs reais sem segmentação estrita de tarefas.
   • Modelagem de Recompensa: Dependendo da tarefa, os sinais de recompensa são derivados de funções de energia contrastiva (treinadas em dados de exploração/demonstração) ou cabeças de classificação leves que identificam comportamentos relevantes para a tarefa.
   • Otimização: A política é otimizada inteiramente dentro do modelo de mundo usando o algoritmo de Otimização de Política Baseada em Modelo Desacoplada (DMO). O modelo local é continuamente ajustado (fine-tuned) durante o treinamento de RL para manter a precisão em relação à distribuição de ações da política atual.

Principais Contribuições

• Framework FoG-MBRL Desacoplado: Os autores introduzem um método que separa a simulação direta de alta fidelidade (via grandes modelos de difusão/fluxo) da diferenciação reversa eficiente (via RSSMs leves). Isso resolve o compromisso entre complexidade de modelagem e custo de computação de gradiente.
• Aprendizado Livre de Simulador a partir de Dados Reais: A abordagem treina agentes inteiramente do zero usando dados de robótica real, contornando a necessidade de simuladores de física e abordando diretamente o gap sim-to-real.
• Implantação Zero-Shot: O framework demonstra a capacidade de treinar políticas no modelo de mundo e implantá-las em robôs reais sem necessidade de ajuste fino adicional ou orientação de clonagem de comportamento (behavior cloning).
• Robustez a Dinâmicas Complexas: O método lida com sucesso com tarefas envolvendo observabilidade parcial, dinâmicas de contato complexas e altos graus de liberdade (ex: manipulação humanoide).

Resultados Experimentais

Os autores avaliaram sua abordagem em três tarefas robóticas do mundo real de complexidade crescente:

1. Push-T: Uma tarefa de manipulação de mesa empurrando um objeto em forma de T.
2. Ego-Centric Push Cube: Um robô quadrúpede empurrando um cubo para um objetivo, envolvendo dinâmicas de caminhada e oclusão parcial.
3. Humanoid Grasp and Lift: Um humanoide de alto DoF agarrando e levantando uma caixa usando uma mão de destreza.

Principais Descobertas:

• Eficiência de Amostragem e de Tempo: O método DMO proposto superou significativamente o PPO (Proximal Policy Optimization) tanto em eficiência de amostragem quanto em tempo de execução (wall-clock time) em todas as tarefas. O PPO teve dificuldade em convergir ou exigiu amostras excessivas, frequentemente caindo em soluções locais subótimas.
• Necessidade de Modelos Globais: Estudos de ablação ("No Diffusion") substituindo o modelo de difusão global por um RSSM padrão para passagens diretas resultaram em falha total da tarefa (0% de taxa de sucesso), confirmando que modelos diretos de alta fidelidade são essenciais para manipulação complexa.
• Superioridade sobre Behavior Cloning: As políticas aprendidas superaram as bases de Behavior Cloning (Action Chunking Transformer), demonstrando a capacidade de descobrir estratégias de longo prazo e lidar com regiões do espaço de estados não bem cobertas pelas demonstrações.
• Sucesso Zero-Shot: Políticas treinadas inteiramente no modelo de mundo alcançaram altas taxas de sucesso (ex: 9/10 em Push-T, 8/10 em Humanoid Grasp) quando implantadas em robôs reais sem treinamento adicional no mundo real.

Significância e Alegações

O artigo alega que aprender agentes inteiramente dentro de modelos de mundo baseados em dados é um caminho viável e promissor para resolver tarefas difíceis de modelar em espaço de imagem. Ao desacoplar os requisitos de precisão direta e tratabilidade reversa, os autores demonstram que modelos de mundo de grande escala e computacionalmente pesados podem ser efetivamente utilizados para a otimização eficiente de políticas.

O trabalho se posiciona como uma solução para a dependência de simuladores de física feitos à mão, oferecendo um método que escala para ambientes complexos e não estruturados através da aquisição de dados do mundo real. Os autores observam que, embora sua implementação atual exija o treinamento de modelos globais do zero e dependa de rótulos de recompensa manuais para certas tarefas, o framework fornece um proxy seguro para o treinamento de robôs reais, permitindo a otimização de políticas escalável sem o risco de danos ao hardware. O trabalho futuro é identificado como o uso de modelos de vídeo pré-treinados em escala de internet para reduzir ainda mais os requisitos de dados e automatizar a rotulagem de recompensas.