Can Tabular Foundation Models Guide Exploration in Robot Policy Learning?

1. Enunciado del Problema

El artículo aborda el desafío de la eficiencia de muestras y la estabilidad de la optimización en el control continuo de alta dimensión para robótica utilizando Aprendizaje por Refuerzo (RL) Profundo.

• Compensación Local vs. Global: Los métodos estándar basados en gradientes (por ejemplo, TD3, SAC) realizan actualizaciones locales en espacios de parámetros de alta dimensión. Aunque son escalables, a menudo se quedan atrapados en regiones subóptimas o convergen lentamente debido a paisajes no convexos. Por el contrario, los métodos de búsqueda global (por ejemplo, Estrategias Evolutivas, Entrenamiento Basado en Poblaciones) exploran regiones más amplias pero incurren en costos de despliegue (interacciones con el entorno) prohibitivos debido a la "maldición de la dimensionalidad".
• Limitaciones de los Modelos Sustitutos: La Optimización Bayesiana (BO) intenta reducir los costos de despliegue utilizando modelos sustitutos. Sin embargo, la BO clásica falla en espacios de alta dimensión (de decenas a cientos de miles de parámetros) debido a la escasez de datos. Además, el paisaje de la política en RL es no estacionario (evoluciona durante el entrenamiento), violando las suposiciones de dominio fijo de la BO tradicional.
• La Brecha: Existe la necesidad de un método que combine la eficiencia de muestras de la búsqueda guiada por sustitutos con la adaptabilidad requerida para políticas de RL de alta dimensión y en evolución, sin requerir presupuestos masivos de interacción.

2. Metodología: TFM-S3

Los autores proponen TFM-S3 (Búsqueda de Subespacio Guiada por Modelo Fundacional Tabular), un marco híbrido que intercala actualizaciones locales estándar basadas en gradientes con rondas periódicas de búsqueda global guiada por sustitutos.

A. Construcción Dinámica de Subespacio de Baja Dimensión

En lugar de buscar en todo el espacio de parámetros de alta dimensión ($\mathbb{R}^D$), TFM-S3 restringe la búsqueda global a un subespacio de baja dimensión ($\mathbb{R}^r$) actualizado dinámicamente.

• Base Informada por Gradientes: En cada ronda de búsqueda $t$, el método recopila las $Q$ instantáneas de gradiente de política más recientes ($G^{(t)}$).
• Descomposición SVD: Realiza una Descomposición en Valores Singulares (SVD) truncada sobre $G^{(t)}$ para extraer los $r$ vectores singulares izquierdos superiores ($U_r$). Estos vectores forman una base ortonormal $A^{(t)}$ que captura las direcciones de actualización dominantes observadas durante el entrenamiento reciente.
• Adaptabilidad: El subespacio evoluciona a medida que avanza el entrenamiento, asegurando que la geometría de la búsqueda se adapte al paisaje de políticas cambiante en lugar de depender de proyecciones aleatorias estáticas.

B. Búsqueda Global Guiada por Sustitutos

Dentro del subespacio de baja dimensión, TFM-S3 realiza una búsqueda iterativa para encontrar candidatos de política de alto rendimiento con despliegues mínimos.

1. Inicialización del Contexto: Se muestrea un pequeño conjunto de políticas candidatas en el subespacio y se evalúan mediante despliegues reales en el entorno para formar un conjunto de contexto ($C$).
2. Modelo Fundacional Tabular (TFM): Se utiliza un Modelo Fundacional Tabular preentrenado (específicamente TabPFN-v2) como sustituto. A diferencia de los Procesos Gaussianos tradicionales, los TFM están preentrenados en diversas tareas de regresión, proporcionando fuertes sesgos inductivos y una generalización robusta a partir de tamaños de muestra muy pequeños (por ejemplo, $K=16$).
3. Bucle de Refinamiento Iterativo:
   • El TFM predice retornos para un gran grupo de nuevos candidatos ($N=256$) muestreados alrededor del mejor candidato actual.
   • Se selecciona el candidato con el retorno predicho más alto.
   • Solo este único mejor candidato es evaluado mediante un despliegue real.
   • El resultado se añade al conjunto de contexto, actualizando la condición del TFM para la siguiente iteración.
   • Este bucle se repite $T$ veces (por ejemplo, 16 iteraciones) por ronda de búsqueda.
4. Actualización de la Política: El mejor candidato encontrado durante la ronda de búsqueda se proyecta de nuevo al espacio original de alta dimensión y se utiliza para inicializar la siguiente fase de entrenamiento local basado en gradientes.

3. Contribuciones Clave

1. Búsqueda de Subespacio Dinámico: Los autores introducen un procedimiento de SVD informado por gradientes para construir subespacios de políticas de baja dimensión que se adaptan a la trayectoria de optimización en evolución, resolviendo el problema de no estacionariedad inherente al RL.
2. Sustitutos de Modelos Fundacionales: Demuestran que los Modelos Fundacionales Tabulares preentrenados pueden servir como sustitutos efectivos y sin ajuste para la predicción de retornos en RL, permitiendo una selección confiable de candidatos con datos extremadamente limitados (generalización con muestras pequeñas).
3. Marco Híbrido Eficiente en Muestras: TFM-S3 conecta exitosamente el RL basado en gradientes y la búsqueda global, logrando una convergencia más rápida y un mejor rendimiento final que las líneas base bajo presupuestos de despliegue idénticos.

4. Resultados Experimentales

El método fue evaluado en tres benchmarks de control continuo MuJoCo: HalfCheetah-v5, Ant-v5 y Humanoid-v5, utilizando TD3 como columna vertebral.

• Rendimiento: TFM-S3 aceleró consistentemente la convergencia en etapas tempranas y logró retornos finales más altos en comparación con TD3 estándar, Búsqueda Aleatoria en el mismo subespacio y una variante "One-Shot" (sin refinamiento iterativo).
• Eficiencia: Bajo un presupuesto fijo de 1 millón de pasos de entorno, TFM-S3 alcanzó umbrales de rendimiento objetivo significativamente más rápido (por ejemplo, alcanzando el 90% del rendimiento final en ~38% de los pasos para Humanoid frente al 48% para TD3 estándar).
• Estabilidad: El método redujo la variabilidad entre semillas aleatorias, indicando actualizaciones de políticas más estables.
• Dinámicas Internas:
  • Consistencia de Clasificación: El TFM mostró una alta correlación de rango de Spearman (hasta 0.97) entre los retornos predichos y reales después de una fase inicial de calentamiento.
  • Calidad de Selección: El candidato top-1 seleccionado por el TFM cayó dentro del top 20% de retornos reales el 63.5% de las veces, en comparación con solo el 20% para la selección aleatoria.
  • Beneficio Iterativo: La variante "One-Shot" tuvo un rendimiento inferior a la versión iterativa completa, demostrando que el refinamiento secuencial (actualizar el sustituto con nuevos datos reales) es crítico para el éxito.

5. Significado

• Nuevo Paradigma para RL: Este trabajo establece a los Modelos Fundacionales Tabulares como herramientas poderosas para el aprendizaje de políticas en robótica, trascendiendo su uso tradicional en datos tabulares estáticos hacia problemas de control dinámicos y de alta dimensión.
• Cerrando la Brecha: Ofrece una solución práctica al cuello de botella de la "eficiencia de muestras" en robótica. Al combinar la reducción de dimensionalidad (SVD) con priores eficientes en datos (TFM), permite la exploración global sin el costo computacional masivo de los métodos basados en poblaciones.
• Modularidad: El marco es ortogonal al algoritmo de RL subyacente (demostrado con TD3 pero aplicable a SAC, PPO, etc.), lo que lo convierte en una mejora versátil para las pipelines existentes de optimización de políticas.

En conclusión, TFM-S3 demuestra que aprovechar modelos fundacionales preentrenados dentro de variedades de baja dimensión reconstruidas dinámicamente puede mejorar drásticamente la eficiencia de muestras y la robustez del aprendizaje de políticas de robots.