Regularized Latent Dynamics Prediction is a Strong Baseline For Behavioral Foundation Models

Resumo Técnico: Regularized Latent Dynamics Prediction (RLDP)

1. O Problema

O artigo aborda o desafio do Aprendizado por Reforço (RL) Zero-Shot e o desenvolvimento de Modelos Fundamentais Comportamentais (BFMs). O objetivo dos BFMs é criar agentes capazes de se adaptar a qualquer recompensa ou tarefa desconhecida sem treinamento adicional, apenas após uma fase de pré-treinamento.

• Limitação Atual: O desempenho dos BFMs depende criticamente da representação de estado aprendida ($\phi$). Métodos atuais de ponta (como FB, PSM, HILP) utilizam objetivos complexos de aprendizado de representação baseados em Medidas Sucessoras (Successor Measures). Esses métodos exigem a estimativa iterativa de backups de Bellman (avaliação ou otimalidade) sobre um conjunto pré-definido de políticas ou recompensas.
• Desafios Específicos:
  1. Instabilidade e Viés: O uso de backups de Bellman em aproximação de funções pode levar a viés, instabilidade e colapso de características (feature collapse).
  2. Cobertura de Dados: Em cenários de baixa cobertura (onde o conjunto de dados offline não cobre todas as ações possíveis), as políticas usadas para calcular os backups podem selecionar ações fora da distribuição (OOD), levando a generalizações incorretas e representações degeneradas.
  3. Complexidade: Os objetivos atuais são computacionalmente caros e sensíveis a hiperparâmetros.

A questão central investigada é: Objetivos complexos de aprendizado de representação são realmente necessários para o RL Zero-Shot, ou uma abordagem mais simples baseada em dinâmica latente é suficiente?

2. Metodologia: RLDP

Os autores propõem o RLDP (Regularized Latent Dynamics Prediction), uma abordagem simplificada que substitui a estimativa complexa de medidas sucessoras por uma previsão de dinâmica latente com regularização.

• Previsão de Dinâmica Latente:

  • O método aprende um encoder de estado $\phi: S \to \mathbb{R}^d$ e um modelo de dinâmica latente $g$ que prevê o próximo estado latente dado o estado atual e a ação: $\phi(s') \approx g(\phi(s), a)$.
  • A função de perda principal ($L_d$) minimiza o erro quadrático entre a previsão da dinâmica e o estado latente real codificado.
  • Vantagem: Diferente dos métodos baseados em Bellman, a previsão de dinâmica é independente da política, evitando problemas de ações OOD durante o treinamento da representação.

• O Problema do Colapso de Características:

  • Os autores observaram que, na sua forma ingênua, a previsão de dinâmica latente leva a um "colapso suave" (mild feature collapse). O modelo tende a mapear estados diferentes para representações latentes muito similares (alta similaridade cosseno) para minimizar a perda, reduzindo a capacidade de distinguir tarefas diferentes.

• Solução: Regularização de Ortogonalidade:

  • Para mitigar o colapso, o RLDP introduz uma regularização de ortogonalidade ($L_r$).
  • As representações de estado são projetadas em uma hipersfera ($S^{d-1}$) e a perda penaliza a similaridade cosseno entre quaisquer dois estados diferentes.
  • Função de Perda Final: $L_{RLDP} = L_d + \lambda L_r$, onde $\lambda$ controla a força da regularização. Isso força a diversidade das características, mantendo o "span" (espaço de cobertura) das funções de recompensa que o modelo pode representar.

• Fase de Inferência (Zero-Shot):

  • Após o pré-treinamento das representações (com dados offline sem recompensa), o BFM é treinado alternando entre a estimativa de medidas sucessoras (usando as representações congeladas de RLDP) e a melhoria da política.
  • Para uma nova tarefa com recompensa $r_{test}$, o vetor de tarefa $z$ é inferido via regressão linear ($r_{test} \approx \phi(s)^\top z$) e a política ótima é extraída.

3. Contribuições Principais

1. Proposta do RLDP: Introdução de uma nova função de perda simples para aprendizado de características de estado para BFMs, combinando previsão de dinâmica latente com regularização de ortogonalidade.
2. Identificação e Mitigação do Colapso: Demonstração de que a previsão de dinâmica latente pura sofre de colapso de características e que uma regularização ortogonal simples é suficiente para resolver isso, mantendo a diversidade das representações.
3. Desempenho Competitivo e Robustez: Evidência empírica de que o RLDP iguala ou supera métodos complexos de ponta (FB, PSM, HILP) em diversos domínios (locomoção, navegação, humanoides).
4. Superioridade em Baixa Cobertura: Demonstração de que o RLDP é robusto em cenários de baixa cobertura de dados, onde métodos baseados em backups de Bellman falham devido a erros de generalização fora da distribuição.

4. Resultados Experimentais

Os experimentos foram conduzidos em três eixos principais:

• RL Zero-Shot Offline (DeepMind Control Suite):

  • Comparado contra Random Features, FB, PSM e HILP em tarefas como Walker, Cheetah, Quadruped e Pointmass.
  • Resultado: O RLDP alcançou desempenho competitivo, frequentemente superando ou empatando com os métodos mais complexos (PSM e FB), utilizando a mesma dimensão de representação (512) e orçamento de gradientes.

• RL Zero-Shot Online (Ambiente Humanoide SMPL):

  • Testado em um ambiente complexo de 358 dimensões de observação e 69 de ação, utilizando o método CPR (Conditional Policy Regularization) para treinamento online.
  • Resultado: O RLDP foi competitivo com as linhas de base, superando-as em tarefas específicas (como "raisearms" e "lieonground") e demonstrando capacidade de generalização em tarefas complexas de controle corporal.

• Cenários de Baixa Cobertura (D4RL - MuJoCo):

  • Avaliação em conjuntos de dados "medium" e "medium-expert" (Hopper, HalfCheetah, Walker2D).
  • Resultado: O RLDP superou consistentemente todas as linhas de base em 5 de 6 tarefas. Em contraste, métodos que dependem de backups de Bellman (como FB) sofreram degradação significativa de desempenho devido à exploração de ações OOD.

• Estudos de Ablação:

  • A regularização de ortogonalidade foi provada essencial: sem ela ($\lambda=0$), a similaridade entre estados aumenta e o desempenho cai.
  • A normalização esférica (Spherical Normalization) no encoder de dinâmica também foi crucial para a estabilidade e desempenho.

5. Significado e Impacto

Este trabalho é significativo porque desafia a suposição de que objetivos complexos de aprendizado de representação são necessários para BFMs eficazes.

• Simplicidade e Eficiência: O RLDP oferece uma alternativa escalável e mais simples, eliminando a necessidade de backups de Bellman iterativos durante o aprendizado de representações, o que facilita a otimização e reduz a complexidade computacional.
• Robustez Prática: Ao evitar a dependência de políticas específicas durante o treinamento de características, o RLDP é mais robusto a conjuntos de dados offline com cobertura limitada, um cenário comum em aplicações do mundo real (como robótica).
• Fundação para Futuros Trabalhos: O sucesso do RLDP sugere que a diversidade de características pode ser garantida de forma mais direta do que se pensava, abrindo caminho para métodos de RL Zero-Shot mais acessíveis e aplicáveis a sistemas com grandes espaços de estado-ação.

Em resumo, o papel demonstra que a previsão de dinâmica latente regularizada é uma base forte e prática para construir Modelos Fundamentais Comportamentais, superando a complexidade desnecessária dos métodos atuais sem sacrificar o desempenho.