Causal Imitation Learning Under Measurement Error and Distribution Shift

1. Planteamiento del problema

El artículo aborda el aprendizaje de imitación offline (IL) en entornos donde coexisten dos desafíos críticos:

1. Error de medición: El estado real relevante para la decisión no se observa directamente. En su lugar, el aprendiz observa proxies ruidosos (mediciones) de un componente de estado latente.
2. Desplazamiento de distribución (Distribution Shift): La distribución conjunta del estado latente y el proceso de medición puede cambiar entre el dominio de entrenamiento (donde se recolectan las demostraciones del experto) y el dominio de despliegue.

En tales regímenes, el Clonación de Comportamiento (BC) estándar falla. El BC típicamente aprende una política $\pi(a|o)$ minimizando la probabilidad condicional observacional de las acciones dadas las observaciones. Los autores argumentan que, bajo un desplazamiento de distribución, estas correlaciones observacionales se vuelven espurias. Específicamente:

• Condicionar sobre mediciones puras ($\pi_{BC2}$) hace que la política sea frágil ante cambios en el mecanismo de medición (por ejemplo, cambios en la calibración de los sensores).
• Ignorar las mediciones y condicionar solo sobre los estados observados ($\pi_{BC1}$) hace que la política sea frágil ante cambios en la dinámica latente o en los núcleos de transición de estado.
• Crucialmente, este sesgo es sistemático y no desaparece incluso con datos infinitos o trayectorias más largas.

2. Metodología

A. Marco Causal y Política Objetivo

Los autores proponen una nueva política objetivo, CausIL, definida no por el condicionamiento observacional, sino por la intervención contrafactual.

• Modelo: El sistema consiste en un estado observado $S_t$, un estado latente $U_t$ (no observado), una medición ruidosa $W_t$ (proxy de $U_t$) y una acción del experto $A_t$. El experto actúa basándose en $S_t$ y un estado latente retrasado $U_{t-1}$ (es decir, $A_t \sim \pi_E(S_t, U_{t-1})$).
• Definición del Objetivo: En lugar de imitar $p(A_t | S_t, W_{t-1})$, el objetivo es aprender la política causal-óptima $\pi_{opt}(s)$, que maximiza la probabilidad de la acción del experto bajo la intervención $do(S_t = s)$:
  $$\pi_{opt}(s) := \arg\max_{a \in \mathcal{A}} p(A_t^{(s)} = a)$$
  donde $A_t^{(s)}$ es el resultado potencial de la acción si el estado fuera fijado en $s$.
• Robustez: Este objetivo depende únicamente de la distribución marginal del estado latente $U_{t-1}$ y del mecanismo invariante del experto $\pi_E$. Es robusto a cambios en el canal de medición $p(W_{t-1}|U_{t-1})$ y, bajo ciertas condiciones, robusto a cambios en los núcleos de transición de estado que preserven la distribución marginal de $U_{t-1}$.

B. Identificación

Dado que $U_{t-1}$ no es observado, los autores derivan las condiciones bajo las cuales $\pi_{opt}$ es identificable únicamente a partir de datos observacionales. Adaptan el marco de Inferencia Causal Proximal (PCI), tratando el estado observado rezagado $Z_t = S_{t-1}$ como un "proxy inductor de tratamiento" y la medición ruidosa $W_{t-1}$ como un "proxy inductor de resultado" para el confundidor latente $U_{t-1}$.

• Caso Discreto: La identificación se logra mediante la inversión de matrices. Si los proxies $Z_t$ y $W_{t-1}$ son suficientemente informativos (satisfaciendo una condición de rango), la distribución intervencional se identifica como:
  $$p(A_t^{(s)}) = P_{A|Z',s} (P_{W'|Z',s})^{-1} P_{W'}$$
  donde las matrices $P$ representan las probabilidades condicionales y marginales de los proxies.
• Caso Continuo: La identificación se basa en las suposiciones de completitud y puente de confusión (confounding bridge). Los autores muestran que existe una "función puente" $h_a(w, s)$ que satisface una ecuación integral de Fredholm de la primera especie, vinculando los datos observados con la distribución contrafactual:
  $$p(A_t^{(s)} = a) = \int h_a(w, s) p(w) dw$$

C. Estimación

El artículo propone estimadores tanto para entornos discretos como continuos:

1. Estimador Discreto: Un estimador de inserción (plug-in) utilizando frecuencias empíricas para construir las matrices de probabilidad y realizar la inversión de matrices derivada en el teorema de identificación.
2. Estimador Continuo: Un procedimiento de estimación adversarial sobre Espacios de Hilbert de Núcleo Reproductor (RKHS). La función puente $h_a$ se aprende resolviendo un problema minimax regularizado (optimización de punto de silla) que impone la restricción de momento condicional $E[Y - h(W, S) | Z, S] = 0$.

3. Contribuciones Clave

• Formulación Teórica: El artículo define formalmente el régimen donde coexisten el error de medición y el desplazamiento de distribución, y demuestra que los objetivos de BC estándar están fundamentalmente sesgados en este escenario.
• Objetivo Causal: Introduce $\pi_{opt}$, una política objetivo que es robusta a cambios en el proceso de medición y en la dinámica latente, siempre que el mecanismo del experto sea invariante.
• Resultados de Identificación: Proporciona condiciones de identificación novedosas para el aprendizaje de imitación offline con confundidores no observados y error de medición, aprovechando las ideas de la inferencia causal proximal (específicamente el uso de estados rezagados y mediciones ruidosas como proxies).
• Algoritmos: Desarrolla estrategias de estimación prácticas tanto para espacios de estado discretos como continuos, incluyendo una solución de forma cerrada para el caso discreto y un estimador adversarial basado en RKHS para el caso continuo.

4. Resultados Experimentales

Los autores evalúan CausIL contra las líneas base estándar de BC ($\pi_{BC1}$ y $\pi_{BC2}$) en dos entornos:

1. Datos Longitudinales Semi-Simulados (PhysioNet 2019):

   • Configuración: Utilizando signos vitales de la UCI como estados observados ($S_t$) y niveles de lactato como un contexto latente retrasado ($U_t$). Las acciones del experto (intensidad del soporte hemodinámico) se simulan basadas en $S_t$ y $U_{t-1}$.
   • Desplazamientos: El estudio introduce "desplazamientos de población" (diferentes composiciones de trayectoria) y "desplazamientos de medición" (cambios en el ruido/sesgo del proxy de lactato).
   • Resultado: CausIL alcanzó consistentemente un Error Cuadrático Medio (MSE) más bajo y se mantuvo estable ante los desplazamientos. En contraste, $\pi_{BC2}$ se degradó drásticamente bajo desplazamientos de medición, y $\pi_{BC1}$ mostró un aumento de error bajo desplazamientos de población.

2. Estudio de Simulación:

   • Configuración: Datos sintéticos generados a partir del modelo de la Sección 2.
   • Resultado: Los experimentos confirmaron las afirmaciones de robustez teórica. $\pi_{opt}$ permaneció invariante cuando el canal de medición cambió, mientras que $\pi_{BC2}$ falló. Del mismo modo, $\pi_{opt}$ fue más estable que $\pi_{BC1}$ cuando los núcleos de transición de estado cambiaron, pero la distribución marginal latente permaneció fija.

5. Significado y Reivindicaciones

El artículo afirma que CausIL ofrece una solución fundamentada al problema de la "ilusión causal" en el aprendizaje de imitación offline cuando los datos son ruidosos y los entornos cambian.

• Robustez: La principal significancia es la demostración de que apuntar a cantidades intervencionales (vía identificación proximal) produce políticas que generalizan mejor a nuevos dominios que los objetivos observacionales, incluso sin señales de recompensa o consultas interactivas al experto.
• Practicidad: El método funciona en entornos totalmente offline, lo cual es crítico para dominios de alto riesgo como la atención médica (ej. soporte de decisiones en UCI) y la robótica, donde la exploración interactiva es imposible o peligrosa.
• Limitaciones: El enfoque depende de suposiciones de identificación específicas (ej. la existencia de proxies válidos y las condiciones de rango/completitud). Los autores señalan que estas condiciones no son triviales, pero se cumplen en muchos escenarios prácticos que involucran estados latentes retrasados y sensores ruidosos.

Los autores concluyen que en entornos realistas y cambiantes con error de medición, aprender políticas a partir de demostraciones puramente observacionales requiere ir más allá de la clonación de comportamiento hacia marcos causales que modelen explícitamente la relación entre estados latentes, proxies y acciones.