Resumen Técnico: ConTraIRL – Abstracciones Contrastivas Factorizadas para la Transferencia de IRL

1. Planteamiento del Problema

El Aprendizaje por Refuerzo Inverso (IRL, por sus siglas en inglés) tiene como objetivo recuperar la función de recompensa subyacente al comportamiento experto, permitiendo que los agentes aprendan habilidades complejas sin necesidad de un diseño manual de recompensas. Un desafío crítico en este dominio es la generalización compositiva: la capacidad de transferir una recompensa aprendida a entornos objetivo donde el contexto es una combinación inédita de dinámicas ambientales ($d$) y objetivos de tarea ($g$) observadas previamente.

Los métodos estándar de IRL suelen modelar la recompensa como una función monolítica del estado. Este enfoque entrelaza características dependientes de la dinámica con aquellas dependientes del objetivo. En consecuencia, cuando un agente se despliega en un entorno objetivo donde el contexto es una combinación $(d_{tar}, g_{tar})$ que nunca se observó durante el entrenamiento, aunque $d$ y $g$ sean individualmente conocidos, la recompensa aprendida se consulta bajo un contexto fuera de la distribución (out-of-distribution). Esto conduce a una generalización poco fiable y a una recuperación deficiente de la recompensa.

El artículo aborda esto estudiando un escenario donde:

• El conjunto de datos de entrenamiento consiste en múltiples entornos donde cada factor ($d$ y $g$) aparece en varios contextos.
• Un subconjunto de combinaciones de factores se reserva como contextos objetivo.
• En los contextos objetivo, solo se dispone de un subconjunto fijo de estados de una única trayectoria experta (supervisión de pocos disparos o few-shot), en lugar de demostraciones completas.

2. Metodología: ConTraIRL

Los autores proponen ConTraIRL (Abstracciones Contrastivas para la Transferencia de IRL), un marco diseñado para permitir la transferencia de recompensas compositivas mediante el aprendizaje de representaciones latentes desacopladas de las dinámicas y los objetivos.

2.1 Abstracciones Contrastivas Factorizadas

El núcleo del marco es una arquitectura de doble codificador que mapea las observaciones en dos espacios latentes separados:

1. Espacio de Dinámicas ($Z_d$): Una representación latente invariante al objetivo $g$.
2. Espacio de Objetivos ($Z_g$): Una representación latente invariante a la dinámica $d$.

Esta factorización asegura que la recompensa recuperada no codifique correlaciones espurias inducidas por la distribución conjunta de entrenamiento. Para lograrlo, los autores imponen la Ortogonalidad de Factores:

• $z_d = \Phi_d(s, d)$ debe ser invariante a $g$.
• $z_g = \Phi_g(s, g)$ debe ser invariante a $d$.

2.2 Alineación de Fase Temporal

Los autores argumentan que la factorización por sí sola es insuficiente para tareas de control secuencial, ya que los estados a lo largo de una trayectoria representan diferentes etapas de ejecución. ConTraIRL introduce la Alineación de Fase Temporal, requiriendo que los estados que corresponden a un progreso relativo similar dentro de un comportamiento se mapeen a coordenadas latentes cercanas a través de las trayectorias.

• Los estados con las mismas dinámicas y una fase similar permanecen cerca en $Z_d$, independientemente del objetivo.
• Los estados con el mismo objetivo y una fase similar permanecen cerca en $Z_g$, independientemente de las dinámicas.
• La fase $\psi(s)$ se define como el paso de tiempo normalizado $t/T$.

2.3 Recuperación de la Recompensa mediante Similitud Latente

La función de recompensa recuperada $R(s, d, g)$ se define como la similitud promedio entre la representación latente del estado actual y las abstracciones expertas condicionadas por la fase:
$$R(s, d, g) = \frac{1}{2} \left( \text{sim}(\Phi_d(s, d), \bar{z}_d(\psi(s))) + \text{sim}(\Phi_g(s, g), \bar{z}_g(\psi(s))) \right)$$
Aquí, $\bar{z}_d(\cdot)$ y $\bar{z}_g(\cdot)$ son los centroides de los estados expertos agrupados por su fase temporal relativa. Esta formulación convierte una combinación fuera de la distribución en una recombinación dentro de la distribución de los factores.

2.4 Objetivos de Entrenamiento

El marco emplea un objetivo contrastivo dual:

1. Aprendizaje de Estructura del Experto ($L_E$): Una pérdida contrastiva basada en márgenes aplicada a pares de expertos para estructurar los espacios latentes. Los pares que comparten el mismo factor y una fase similar se tratan como positivos; los demás, como negativos. Esto moldea los espacios de los expertos de forma independiente al aprendiz.
2. Calibración de Experto-Aprendiz ($L_L$): Un objetivo de contraste de ruido que calibra la recompensa distinguiendo entre estados del experto ($x_E$) y estados del aprendiz ($x_L$). Esto minimiza efectivamente la divergencia de Jensen–Shannon entre las ocupaciones de estado del experto y del aprendiz.

El objetivo total es $L_{total} = L_E - L_L$.

2.5 Detalles de Implementación

• Codificador Modulado por Contexto: En lugar de una simple concatenación, el marco utiliza una capa de entrada inspirada en redes hiperredes (hypernetworks) donde el descriptor de contexto ($d$ o $g$) genera matrices de pesos y vectores de sesgo específicos para transformar el estado. Esto influye dinámicamente en la extracción de características, fomentando el aislamiento de la estructura relevante del factor en las etapas tempranas de la red.
• Búfer de Abstracción del Experto: Un búfer mantiene los centroides del experto condicionados por la fase, los cuales se actualizan tras cada actualización del codificador para asegurar la consistencia con la evolución del espacio latente.

3. Contribuciones Clave

• Representación Factorizada: La propuesta de una arquitectura de doble codificador que desacopla explícitamente las dinámicas y los objetivos en espacios latentes ortogonales, evitando el entrelazamiento de características que afecta a los modelos de IRL monolíticos.
• Alineación de Fase Temporal: La integración de la fase temporal relativa en el criterio de alineación latente, proporcionando una guía densa para la recuperación de la recompensa en tareas secuenciales en lugar de depender únicamente de la coincidencia de estados dispersos.
• Transferencia Compositiva con Supervisión de Pocos Disparos: Un marco capaz de recuperar recompensas para combinaciones inéditas de dinámicas y objetivos utilizando solo un subconjunto parcial de estados expertos del entorno objetivo.
• Aprendizaje de Estructura del Experto: Un objetivo explícito para dar forma a los espacios de los expertos de forma independiente a la distribución del aprendiz, proporcionando una referencia estacionaria para la recuperación de la recompensa.

4. Resultados Experimentales

Los autores evaluaron ConTraIRL en cuatro entornos de control continuo de MuJoCo (Ant, HalfCheetah, Walker, Swimmer) con variaciones compositivas en las dinámicas y los objetivos.

• Rendimiento en Combinaciones Inéditas: ConTraIRL superó consistentemente a los modelos de referencia (TraIRL, C-AIRL y SFM) en los contextos objetivo. Por ejemplo, en el entorno Ant, ConTraIRL alcanzó un retorno normalizado de 0.93 ± 0.02, frente al 0.88 de SFM y el 0.57 de TraIRL.
• Robustez a la Escasez de Datos: En estudios de ablación donde se redujo el número de estados expertos objetivo (del 20% al 5%), los métodos de referencia sufrieron caídas catastróficas de rendimiento (por ejemplo, TraIRL cayó a 0.05), mientras que ConTraIRL mantuvo un rendimiento significativamente más alto (0.27 con el 5% de estados), demostrando una transferencia efectiva de pocos disparos (few-shot).
• Análisis de Ablación:
  • Eliminar la factorización ("w/o F") causó una caída severa en el rendimiento, confirmando la necesidad de representaciones desacopladas.
  • Eliminar la alineación de fase temporal ("w/o TA") redujo el rendimiento, indicando que la alineación de fase proporciona una guía densa esencial más allá de la simple factorización.
• Robustez al Ruido de Etiquetas: ConTraIRL demostró una mayor robustez al ruido en las etiquetas de los factores en comparación con los modelos de referencia, debido a su dependencia de restricciones de similitud relativa en lugar de la predicción supervisada directa.
• Visualización del Espacio Latente: Las visualizaciones de UMAP confirmaron que las abstracciones aprendidas separaron con éxito las muestras según su factor correspondiente (dinámica u objetivo) mientras permanecían invariantes al otro, y que la fase temporal se preservó como un gradiente suave dentro de los grupos (clusters).

5. Significación y Reivindicaciones

El artículo afirma que ConTraIRL aborda una limitación fundamental en la transferencia de IRL: la incapacidad de generalizar a combinaciones no vistas de factores contextuales conocidos. Al factorizar el contexto en componentes ortogonales de dinámica y objetivo, y al alinearlos temporalmente, el marco permite la generalización compositiva.

Los autores enfatizan que su enfoque permite una recuperación de recompensa fiable en entornos objetivo donde la combinación específica de dinámica y objetivo no estuvo presente durante el entrenamiento, siempre que los factores individuales hayan sido observados. Se demuestra que la recompensa aprendida sigue de cerca el rendimiento real de la tarea, sirviendo como un proxy fiable para la optimía de políticas incluso cuando las recompensas de verdad absoluta (ground-truth) no están disponibles.

El trabajo está delimitado modestamente a la transferencia compositiva a través de factores contextuales conocidos con supervisión limitada en el objetivo. Los autores señalan explícitamente que la formulación actual asume el acceso a etiquetas de factores y fases temporales comparables, e identifican la extensión del enfoque hacia escenarios con factores latentes, comportamientos de longitud variable y transferencia totalmente zero-shot como direcciones importantes para trabajos futuros.