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

1. El Problema: Limitaciones de los Modelos Fundamentales de Comportamiento (BFMs)

Los Modelos Fundamentales de Comportamiento (BFMs) son agentes diseñados para el aprendizaje por refuerzo (RL) zero-shot (sin entrenamiento adicional). Su objetivo es aprender una política óptima para cualquier función de recompensa desconocida, basándose únicamente en un conjunto de datos de interacciones sin recompensas (reward-free).

El funcionamiento de un BFM depende críticamente de dos componentes:

1. Una representación de estado ($\phi$) que mapea el estado $s$ a un espacio latente.
2. Una política condicional a un vector latente $z$ (embedding de tarea), donde la recompensa se asume lineal en el espacio de características: $r(s) = \phi(s)^\top z$.

El desafío principal: La expresividad del BFM está limitada por el "rango" (span) de las funciones de recompensa que pueden representarse con las características aprendidas.

• Los métodos actuales (como Forward-Backward o PSM) utilizan objetivos complejos basados en la estimación de medidas sucesoras (successor measures), que requieren iterar sobre backups de Bellman.
• Estos métodos sufren de inestabilidad, sesgos y, crucialmente, requieren una cobertura de datos suficiente. En escenarios de baja cobertura (donde los datos no cubren bien el espacio de estados-acciones), los backups de Bellman pueden seleccionar acciones fuera de distribución (OOD), llevando a representaciones degeneradas y fallos en la generalización.
• Además, los objetivos de predicción de dinámica latente simples (sin regularización) tienden a sufrir de un colapso de características (feature collapse), donde las representaciones de diferentes estados se vuelven demasiado similares, reduciendo la capacidad del modelo para distinguir tareas.

2. Metodología: RLDP (Predictión de Dinámica Latente Regularizada)

Los autores proponen RLDP, un enfoque simple pero efectivo que evita la complejidad de los backups de Bellman y mitiga el colapso de características mediante una regularización de ortogonalidad.

A. Aprendizaje de Representaciones

En lugar de estimar medidas sucesoras directamente, RLDP aprende una representación de estado $\phi$ mediante la predicción de la dinámica latente:

1. Objetivo Principal: Predecir el siguiente estado latente dado el estado actual y la acción: $\phi(s') \approx g(\phi(s), a)^\top w$.
2. Pérdida de Dinámica ($L_d$): Minimiza el error cuadrático entre la predicción del estado futuro latente y el estado latente real codificado.
3. El Problema del Colapso: Se observó que minimizar solo $L_d$ lleva a una similitud creciente entre las representaciones de diferentes estados (colapso suave), reduciendo el rango de recompensas que el modelo puede resolver.
4. Solución: Regularización de Ortogonalidad ($L_r$): Para mantener la diversidad de características, se añade una pérdida que minimiza la similitud coseno entre cualquier par de estados en el espacio latente. Se proyectan las representaciones en una hiperesfera ($S^{d-1}$) y se penaliza la alineación entre ellas.
   • Pérdida Total: $L_{RLDP} = L_d + \lambda L_r$, donde $\lambda$ controla la fuerza de la regularización.

B. Aprendizaje del BFM (Zero-Shot RL)

Una vez entrenada la representación $\phi$ (que se mantiene congelada):

1. Se entrena un modelo de medida sucesora ($\psi$) y una política ($\pi_z$) utilizando los datos offline.
2. Se utiliza la estructura de características sucesoras para calcular valores de acción para cualquier nueva recompensa $r_{test}$ encontrando el vector $z$ tal que $r_{test}(s) \approx \phi(s)^\top z$.
3. La política se optimiza para maximizar el valor esperado dado este $z$.

3. Contribuciones Clave

1. Propuesta de RLDP: Introducen un objetivo de aprendizaje de representaciones simple basado en la predicción de dinámica latente con regularización de ortogonalidad, eliminando la necesidad de complejos backups de Bellman durante la fase de pre-entrenamiento.
2. Identificación y Mitigación del Colapso: Demuestran empíricamente que la predicción de dinámica latente pura sufre de colapso de características y proponen la regularización de ortogonalidad como una solución efectiva y escalable.
3. Rendimiento en Escenarios de Baja Cobertura: Muestran que RLDP es robusto en conjuntos de datos con baja cobertura (donde otros métodos fallan debido a la generalización incorrecta de acciones OOD), superando a métodos basados en medidas sucesoras en benchmarks como D4RL.
4. Competitividad General: Logran un rendimiento comparable o superior a los métodos state-of-the-art (FB, PSM, HILP) en tareas de control continuo (DeepMind Control Suite) y en entornos humanoides complejos (SMPL), tanto en configuraciones offline como online.

4. Resultados Experimentales

Los experimentos se realizaron en múltiples dominios y configuraciones:

• Control Continuo (DeepMind Control Suite):

  • En tareas offline (Walker, Cheetah, Quadruped, Pointmass), RLDP iguala o supera a FB y PSM.
  • En tareas online con un humanoide complejo (SMPL, 358 dimensiones de observación), RLDP compite favorablemente con el oracle (FB-CPR) en la mayoría de las 45 tareas, superando a otros métodos en tareas específicas como "raisearms" y "lieonground".

• Escenarios de Baja Cobertura (D4RL):

  • En datasets medium y medium-expert de MuJoCo (Hopper, HalfCheetah, Walker2D), RLDP supera significativamente a los baselines en 5 de 6 tareas.
  • Esto confirma que evitar los backups de Bellman durante el aprendizaje de la representación es crucial cuando los datos no cubren todo el espacio de políticas.

• Estudios de Ablación:

  • Regularización de Ortogonalidad: Sin ella ($\lambda=0$), el rendimiento cae drásticamente debido al colapso de características. Con $\lambda > 0$, la diversidad se mantiene y el rendimiento mejora.
  • Normalización Esférica: Proyectar las salidas en una hiperesfera es vital para la estabilidad y el rendimiento.
  • Horizonte de Codificación: El rendimiento es robusto a diferentes horizontes de predicción ($H=1$ a $H=20$), aunque $H=5$ se utilizó como estándar.

5. Significado e Impacto

Este trabajo es significativo porque desafía la noción de que se necesitan objetivos de aprendizaje de representación extremadamente complejos (basados en medidas sucesoras y backups de Bellman) para lograr un RL zero-shot efectivo.

• Simplicidad y Estabilidad: RLDP ofrece un enfoque más simple, estable y fácil de optimizar, que no requiere reconstrucción de estados ni estimación explícita de medidas sucesoras durante la fase de aprendizaje de características.
• Robustez en Datos Reales: Al ser independiente de la política durante el aprendizaje de la representación, RLDP es inherentemente más robusto a la falta de cobertura en los datos, un problema común en aplicaciones del mundo real (robótica) donde recopilar datos exhaustivos es costoso o peligroso.
• Fundamento para Modelos Generales: Proporciona una base sólida y escalable para el desarrollo de Modelos Fundamentales de Comportamiento que pueden adaptarse a nuevas tareas y recompensas sin reentrenamiento, facilitando su aplicación en robots de servicio y sistemas autónomos complejos.

En resumen, el papel demuestra que la predicción de dinámica latente, cuando se combina con una simple regularización de ortogonalidad, es un método potente y subestimado para aprender representaciones universales en RL, superando a enfoques más complejos en escenarios críticos de baja cobertura.