HCLSM: Hierarchical Causal Latent State Machines for Object-Centric World Modeling

Resumen Técnico: HCLSM

1. El Problema

Los modelos del mundo actuales, que predicen estados futuros a partir de video, sufren de limitaciones fundamentales debido a sus representaciones latentes "planas" (flat). Estas representaciones:

• Entrelazan objetos: No distinguen entre entidades discretas (ej. una taza, una mesa, un robot), mezclando sus estados en un único vector.
• Ignoran la estructura causal: No modelan explícitamente qué causa qué, lo que impide el razonamiento contrafactual (ej. "¿qué pasaría si el robot empujara más fuerte?").
• Colapsan la dinámica temporal: Tratan todas las escalas de tiempo (movimiento continuo de milisegundos, eventos discretos como colisiones, y planes estratégicos a largo plazo) con un único mecanismo, sin jerarquía.

Los modelos existentes (como V-JEPA o DreamerV3) logran predecir futuros latentes, pero carecen de la estructura necesaria para la planificación y el razonamiento basado en objetos.

2. Metodología: Arquitectura HCLSM

HCLSM propone una arquitectura de cinco capas interconectadas que opera bajo tres principios: descomposición centrada en objetos, dinámica temporal jerárquica y aprendizaje de estructura causal.

A. Descomposición Centrada en Objetos (Capa 2)

• Slot Attention Dinámico: Utiliza un mecanismo de atención iterativa donde "slots" (ranuras) compiten por tokens de parches de imagen. Incluye una cabeza de existencia para predecir qué slots están activos y un mecanismo de "nacimiento" de slots basado en residuos de atención.
• Decodificador de Difusión Espacial (SBD): Cada slot se proyecta en una cuadrícula espacial y se decodifica mediante una CNN. A diferencia de métodos anteriores que reconstruyen píxeles, HCLSM reconstruye características de un ViT congelado (siguiendo DINOSAUR), lo que proporciona una señal semántica para la descomposición.
• Gráfico de Relaciones: Un GNN procesa las interacciones entre pares de slots, generando pesos de borde que codifican la estructura de interacción.

B. Dinámica Temporal Jerárquica (Capa 3)
El modelo utiliza un motor de tres niveles para manejar diferentes escalas de tiempo:

1. Nivel 0 (SSM Selectivo): Modela la física continua frame-a-frame para cada objeto individualmente usando State Space Models (SSM) selectivos (tipo Mamba).
2. Nivel 1 (Transformer de Eventos Dispersos): Detecta límites de eventos (ej. colisiones) y procesa solo esos timesteps críticos mediante un Transformer disperso, reduciendo la complejidad computacional.
3. Nivel 2 (Transformer de Compresión de Objetivos): Resume la secuencia de eventos en tokens de estado abstracto para el razonamiento de alto nivel y la planificación.

C. Estructura Causal (Capa 4)

• Aprende una matriz de adyacencia causal mediante pesos de GNN y una regularización estilo NOTEARS (restricción de DAG) para asegurar que las interacciones sean acíclicas y causales.

D. Protocolo de Entrenamiento en Dos Etapas (Clave)
El artículo identifica que si se activan todas las pérdidas desde el inicio, el modelo opta por una representación distribuida fácil de predecir en lugar de especializarse por objetos. Para evitar esto, se propone:

• Etapa 1 (40% del entrenamiento): Solo se entrena con la pérdida de reconstrucción SBD y un regularizador de diversidad. Esto fuerza a los slots a especializarse espacialmente (aprender "qué es" cada cosa).
• Etapa 2 (60% restante): Se activa la pérdida de predicción dinámica (estilo JEPA) sobre los objetos ya descompuestos. El modelo aprende "qué hacen" los objetos.

3. Contribuciones Clave

1. Arquitectura Unificada: Primer modelo diferenciable que combina slots centrados en objetos, una jerarquía temporal de tres niveles (SSM + Transformer disperso + Transformer de objetivo) y razonamiento causal basado en GNN.
2. Decodificador SBD con ViT: Uso de características de ViT congeladas como objetivo de reconstrucción para permitir la descomposición de objetos no supervisada en video de robots reales.
3. Protocolo de Entrenamiento en Dos Etapas: Resuelve el conflicto de optimización entre la especialización espacial y la predicción dinámica, imitando el desarrollo del córtex visual humano.
4. Optimización de Rendimiento: Implementación de un kernel personalizado en Triton para el escaneo de SSM, logrando una aceleración de 38x frente a la implementación secuencial de PyTorch.
5. Código Abierto y Datos Reales: Implementación completa (8,478 líneas, 51 módulos) entrenada en el dataset Open X-Embodiment (tarea PushT), con pruebas unitarias y suite de evaluación.

4. Resultados

El modelo se entrenó con 68M de parámetros en la tarea de manipulación robótica PushT (empujar un bloque en forma de T):

• Pérdida de Predicción: 0.008 MSE (Mean Squared Error) para el siguiente estado.
• Descomposición Espacial: La pérdida de reconstrucción SBD alcanzó 0.0075, confirmando que los slots aprendieron a asignarse a regiones espaciales específicas.
• Detección de Eventos: El modelo identificó correctamente 2-3 eventos por secuencia de 16 frames, alineados con cambios de estado significativos (ej. contacto robot-objeto).
• Rendimiento: El kernel Triton redujo el tiempo de predicción temporal por objeto del cuello de botella dominante a solo el 5% del tiempo de paso hacia adelante.
• Comparativa: Sin el protocolo de dos etapas, el modelo lograba una pérdida de predicción menor (0.002) pero sin estructura de objetos (representación distribuida). Con dos etapas, la pérdida es mayor (0.008) pero la estructura es interpretable y útil para la planificación.

5. Significado y Limitaciones

Significado:
HCLSM demuestra que es posible construir modelos del mundo que reflejen la estructura real del entorno (objetos discretos, múltiples escalas de tiempo, causalidad) en lugar de tratar el video como una secuencia de píxeles. El protocolo de dos etapas es una contribución metodológica crucial para estabilizar el aprendizaje de representaciones estructuradas frente a la presión de optimización hacia representaciones planas.

Limitaciones Actuales:

• Conteo de Slots: El modelo no aprende a "matar" slots innecesarios; en una escena de 3 objetos, usa los 32 slots disponibles, dividiendo cada objeto en múltiples slots.
• Descubrimiento Causal Explícito: La matriz de adyacencia causal explícita colapsó a cero bajo regularización de dispersión; la estructura causal actual es implícita a través de los pesos del GNN.
• Estabilidad Numérica: El entrenamiento en precisión bfloat16 es inestable para configuraciones más grandes (Base/Large) y depende fuertemente de la semilla aleatoria (40-60% de las ejecuciones divergen).
• Escala: Actualmente probado solo en una tarea simple (PushT) y con una configuración pequeña (68M).

Futuro:
El trabajo sienta las bases para agentes autónomos completos. Los pasos futuros incluyen inicializar con ViTs preentrenados (V-JEPA 2), entrenar en datasets de manipulación multi-objeto más complejos y cerrar el ciclo con planificadores (CEM/MPPI) para control robótico en bucle cerrado.