Resumo Técnico: Modelo de Mundo Prismático (PRISM-WM)

1. Declaração do Problema

A Aprendizagem por Reforço Baseada em Modelo (MBRL) e o planejamento em domínios robóticos são fundamentalmente limitados pela natureza híbrida da dinâmica física. Sistemas do mundo real, particularmente na locomoção com pernas e manipulação, envolvem movimento contínuo pontuado por eventos discretos (por exemplo, contatos, impactos, aderência versus deslizamento). Essas interações criam dinâmicas não suaves e altamente não lineares com descontinuidades abruptas no campo vetorial subjacente.

Modelos de mundo latente convencionais empregam tipicamente redes neurais monolíticas que impõem continuidade global. Embora teoricamente capazes de aproximação universal, esses modelos aprendem na prática uma representação média que suaviza excessivamente modos dinâmicos distintos. Para planejadores de otimização de trajetória, essa suavização excessiva induz erros cumulativos durante previsões de longo horizonte, particularmente em fronteiras físicas. Isso resulta em:

• Desfoque espectral através das fronteiras de contato, causando incompatibilidade significativa de gradientes.
• Deriva cumulativa em espaços de configuração de alta dimensão, tornando o planejamento de longo prazo pouco confiável.
• Colapso da política em ambientes ricos em contato, onde modelos padrão falham em distinguir entre regimes estáveis e instáveis.

2. Metodologia: O Modelo de Mundo Prismático (PRISM-WM)

Para abordar essas limitações, os autores introduzem o PRISM-WM, uma arquitetura estruturada projetada para decompor dinâmicas híbridas complexas em primitivas composáveis. O modelo é construído sobre uma estrutura de Mistura de Especialistas (MoE) sensível ao contexto com um objetivo inovador de ortonormalização latente.

Arquitetura Central

• Comutação Sensível ao Contexto: Uma rede de comutação ($G$) recebe um vetor de contexto $c_t$ (composto pelo estado $z_t$, ação $a_t$ e, opcionalmente, incorporação de tarefa $e_{task}$) e produz pesos normalizados por softmax $w_t$. Esse mecanismo identifica implicitamente o regime físico ativo (por exemplo, voo versus apoio) sem supervisão explícita.
• Especialistas Especializados: O estado é roteado para um conjunto de $K$ redes especialistas $\{E_k\}$. Cada especialista se especializa em aprender o fluxo contínuo dentro de um regime dinâmico local específico, evitando interferência de modos conflitantes.
• Previsão Residual: O próximo estado latente é previsto como uma combinação ponderada das saídas dos especialistas adicionada ao estado atual:
  $$\hat{z}_{t+1} = z_t + \sum_{k=1}^{K} w_{t,k} \cdot v_k(x_t)$$
  onde $v_k$ representa as dinâmicas residuais previstas pelo $k$-ésimo especialista.

Inovação Chave: Otonormalização Latente

Para prevenir o colapso de modos (onde especialistas aprendem características redundantes) e garantir diversidade, o PRISM-WM introduz uma camada de ortonormalização Gram-Schmidt diferenciável.

• Após cada especialista produzir um vetor da penúltima camada $u_k$, o modelo aplica um processo Gram-Schmidt estabilizado para gerar vetores de base ortogonais $v_k$.
• Isso força as representações dos especialistas a residirem em uma variedade de Stiefel, garantindo que elas abranjam um subespaço de posto máximo e não redundante para as dinâmicas residuais.
• Embora os modos de contato físico sejam desordenados, essa ortonormalização garante que os especialistas capturem componentes distintos da variedade residual, permitindo que a rede de comutação reconstrua modos físicos desordenados por meio de uma soma ponderada.

Integração com Planejamento

O PRISM-WM é integrado ao framework de planejamento TD-MPC. Ao fatorizar dinâmicas complexas em especialistas não redundantes, o modelo reduz o resíduo do modelo de um passo de $\epsilon$ para aproximadamente $\epsilon/K\alpha$. Essa decomposição estrutural mitiga a deriva cumulativa, permitindo rolagens de longo horizonte de alta fidelidade e propagação de gradientes estável para aprendizado de políticas.

3. Contribuições Principais

1. Decomposição Arquitetural: Introdução do PRISM-WM, que decompõe dinâmicas híbridas usando uma arquitetura MoE sensível ao contexto acoplada à ortonormalização latente, afastando-se de aproximações globais monolíticas.
2. Objetivo de Otonormalização: Um objetivo de treinamento inovador usando ortonormalização Gram-Schmidt diferenciável para impor diversidade de especialistas e prevenir colapso de modos, alinhando estruturalmente o espaço latente com transições físicas discretas.
3. Validação Teórica e Empírica: Demonstração de que essa decomposição reduz erros cumulativos em espaços de alta dimensão e melhora a geometria da paisagem de perda tanto para planejamento (TD-MPC) quanto para aprendizado direto de políticas (PWM).
4. Capacidade de Transferência Zero-Shot: Evidência de que a base ortogonal permite que o modelo generalize entre simuladores com físicas de contato diferentes (por exemplo, IsaacGym de contato macio para MuJoCo de contato rígido) isolando modos dinâmicos específicos (por exemplo, atrito versus suporte vertical).

4. Resultados Experimentais

Os autores avaliaram o PRISM-WM em benchmarks de controle contínuo, incluindo humanoides de alta dimensão, configurações de múltiplas tarefas e transferência sim-a-sim zero-shot.

• Locomoção de Alta Dimensão (DiffRL & Humanoid-Bench): O PRISM-WM alcançou maior eficiência de amostragem e desempenho assintótico em comparação com bases monolíticas (TD-MPC2, DreamerV3, PWM, SAC). Notavelmente, enquanto modelos monolíticos sofreram colapso de política em tarefas de locomoção bípede (Humanoid, SNU-Humanoid) devido à média de eventos de contato não suaves, o PRISM-WM manteve otimização estável.
• Generalização Multi-tarefa: Em um benchmark da DeepMind Control Suite com 30 tarefas, o PRISM-WM alcançou uma pontuação normalizada média de 0,531, uma melhoria de 23,5% sobre a base monolítica TD-MPC2 (0,430). Superou a base em 26 de 30 tarefas, mitigando efetivamente a transferência negativa entre dinâmicas heterogêneas.
• Transferência Sim-a-Sim Zero-Shot: Ao transferir do IsaacGym para o MuJoCo, a base MLP monolítica sofreu uma queda de desempenho de 47,8% com uma taxa de sucesso de 12%. O PRISM-WM reduziu essa queda para 11,8% com uma taxa de sucesso de 88%. A ortonormalização permitiu que o modelo ajustasse pesos para especialistas específicos (por exemplo, atrito horizontal) sem comprometer a variedade de dinâmicas global.
• Estabilidade e Eficiência: Estudos de ablação confirmaram que a ortonormalização previne a degradação da estimativa de recompensa observada em modelos MoE não restritos. Testes de eficiência de inferência em uma NVIDIA RTX 4090 mostraram que o PRISM-WM (K=4) alcançou 7.712 FPS, superando ligeiramente a base monolítica (7.499 FPS) com sobrecarga de memória mínima.

5. Significado e Alegações

O artigo afirma que o PRISM-WM fornece um substrato de alta fidelidade para algoritmos de otimização de trajetória, atuando como um modelo fundamental para agentes baseados em modelo. Ao decompor física complexa em primitivas ortogonais, o modelo:

• Mitiga a deriva de longo prazo em horizontes de planejamento, permitindo raciocínio de longo horizonte confiável em ambientes ricos em contato.
• Estabiliza a propagação de gradientes para aprendizado de políticas, transformando paisagens de perda fracturadas em misturas convexas funcionalmente suaves.
• Permite tomada de decisão robusta em sistemas híbridos onde modelos de mundo padrão falham devido à suavização excessiva.

Os autores concluem que decompor física complexa em primitivas ortogonais é "fundamentalmente essencial para dominar o controle rico em contato". Embora o quadro atual exija ajuste manual da contagem de especialistas $K$, o trabalho estabelece uma fundação estrutural para a construção de modelos de mundo de propósito geral e fundamentados na física para futuros agentes incorporados.