TRAM: Test-Time Risk Adaptation with Mixture of Agents

Resumen Técnico: TRAM (Adaptación de Riesgo en Tiempo de Ejecución mediante Mezcla de Agentes)

1. Formulación del Problema

El artículo aborda una brecha crítica en el despliegue del Aprendizaje por Refuerzo (RL): las especificaciones de seguridad a menudo cambian después del entrenamiento. Un agente puede ser entrenado en un conjunto de tareas fuente pero desplegado en un entorno con nuevos mapas de peligros, umbrales de riesgo revisados o diferentes restricciones de alineación conductual (por ejemplo, un nuevo comportamiento de referencia).

Los métodos existentes de RL sensibles al riesgo suelen fijar el modelo de riesgo (por ejemplo, varianza, CVaR o una restricción específica) durante la fase de entrenamiento. Este enfoque es frágil porque:

1. Especialización excesiva: Entrenar con un proxy de riesgo específico (como la varianza de la recompensa) puede colapsar la diversidad de las políticas fuente, eliminando comportamientos que podrían ser seguros bajo una especificación de riesgo diferente y futura.
2. Representación de riesgo inadecuada: Los riesgos basados en la varianza no logran capturar peligros espaciales (por ejemplo, proximidad a una barrera), la deriva conductual respecto a una referencia o la volatilidad local, especialmente en entornos deterministas donde la varianza de la recompensa de la trayectoria es cero.
3. Costos de reentrenamiento: Adaptarse a una nueva especificación de riesgo generalmente requiere reentrenar o ajustar finamente, lo cual es computacionalmente costoso y a menudo inviable una vez que el modelo ha sido desplegado.

El artículo formaliza el problema como adaptación en tiempo de despliegue sin actualizaciones. El objetivo es reutilizar una biblioteca fija de políticas fuente neutrales al riesgo $\{\pi_j\}$ para satisfacer una nueva funcional de riesgo $\rho(\pi)$ especificada en tiempo de prueba sin modificar ningún parámetro fuente. El objetivo es maximizar la recompensa esperada $J(\pi)$ sujeta a una restricción de riesgo $\rho(\pi) \le \delta$, aproximada mediante un sustituto penalizado:
$$\max_{\pi} \tilde{J}_c(\pi) = J(\pi) - c \rho(\pi)$$

2. Metodología: TRAM

El método propuesto, TRAM, desacopla el entrenamiento fuente de la especificación de riesgo. Opera en dos fases:

Fase 1: Entrenamiento Fuente Neutro al Riesgo
Las políticas fuente se entrenan en tareas fuente relacionadas utilizando objetivos estándar neutrales al riesgo. Esto preserva un "casco conductual" diverso, asegurando que la biblioteca contenga una amplia gama de distribuciones de ocupación (patrones de visita estado-acción) que puedan recombinarse posteriormente.

Fase 2: Puntuación y Composición de Fuentes en Tiempo de Ejecución
En el despliegue, TRAM no reentrena. En su lugar, construye una política desplegada $\pi_T$ seleccionando acciones de la biblioteca fuente basándose en una regla de composición puntuada por fuente.

1. Evaluación de Riesgo: Para cada política fuente $\pi_j$, TRAM calcula una puntuación de riesgo global basada en la ocupación, $\rho(d_{\pi_j})$, almacenada en caché. Esto permite funcionales de riesgo flexibles definidos sobre ocupaciones estado-acción con descuento $d_\pi$, tales como:
   • Exposición a Barreras: Tiempo pasado cerca de una región de peligro.
   • Divergencia Conductual: Divergencia KL respecto a un comportamiento de referencia seguro.
   • Volatilidad Local: Fluctuaciones por paso en la ocupación.
2. Estimación de Valor: El valor de la tarea objetivo $Q^{\pi_j}_T(s, a)$ para cada fuente se estima en el estado actual $s$. Esto puede hacerse eficientemente mediante Características Sucesoras (SFs) mediante un producto punto $Q^{\pi_j}_T(s, a) = \psi_{\pi_j}(s, a)^\top w_T$, o mediante simulaciones/críticos.
3. Selección de Acción: La acción desplegada se selecciona maximizando la puntuación de fuente ajustada al riesgo:
   $$S_j(s, a) = Q^{\pi_j}_T(s, a) - c \rho(d_{\pi_j})$$
   $$\pi_T(a|s) \in \arg\max_{b \in A} \max_{j \in [n]} S_j(s, b)$$
   La fuente activa puede cambiar a lo largo de la trayectoria, creando una política "cosida localmente" que responde al estado actual y a la especificación de riesgo específica.

Garantías Teóricas:
El artículo enmarca explícitamente a TRAM como un sustituto puntuado por fuente tratable, no como un optimizador exacto del problema completo de control de ocupación de la política cosida.

• Límites Localizados: El análisis teórico (Teorema 4.1) proporciona límites sobre la suboptimalidad de la política TRAM en relación con la política óptima neutra al riesgo. El límite depende del desajuste de recompensa localizado entre el objetivo y la fuente mejor coincidente, más un término proporcional al peso de la penalización de riesgo $c$ y la constante de Lipschitz del funcional de riesgo.
• Desajuste del Casco Fuente: El artículo reconoce que la ocupación realizada de la política cosida puede diferir del casco convexo de las ocupaciones fuente. El Apéndice C introduce un término medible de "desajuste de costura" ($\epsilon_{stitch}$) que cierra la brecha entre los riesgos fuente almacenados en caché y el riesgo realizado de la política desplegada.

3. Contribuciones Clave

1. Formalización de la Adaptación sin Actualizaciones: El artículo formaliza el escenario de adaptar políticas fuente fijas a riesgos basados en ocupación recién especificados sin actualizaciones de parámetros.
2. Algoritmo TRAM: Introducción de una regla de composición puntuada por fuente que soporta diversos tipos de riesgo (barreras, divergencia, volatilidad) en tiempo de prueba, en contraste con métodos ligados a sustitutos fijos como la varianza.
3. Fundamentos Teóricos:
   • Demostración de que el entrenamiento fuente neutro al riesgo minimiza el arrepentimiento de riesgo de despliegue en el peor de los casos (Teorema 4.3), preservando el casco de ocupación reutilizable más amplio.
   • Garantías de transferencia localizadas que separan la calidad de la transferencia de recompensa de la penalización de riesgo de despliegue.
   • Análisis del "desajuste de costura" entre los riesgos fuente almacenados en caché y los riesgos de la política realizada.
4. Validación Empírica: Demostración en cuatro dominios distintos:
   • Gridworlds: Validación de un modelado de riesgo expresivo (peligros espaciales, volatilidad por paso) donde los métodos basados en varianza fallan.
   • MuJoCo Reacher: Reducción de violaciones de barreras en control continuo manteniendo el rendimiento hacia la meta.
   • Safety-Gymnasium: Logro de probabilidades de impacto con barreras más bajas en comparación con proxies de seguridad congelados y líneas base basadas en varianza.
   • Alineación de LLM: Adaptación de modelos de lenguaje fijos de 7B a nuevas restricciones de riesgo conductual (KL a una referencia segura) sin ajuste fino, mostrando una mejor alineación que la maximización pura de recompensas.

4. Resultados

Los experimentos muestran que TRAM mejora consistentemente la seguridad de despliegue (reduciendo la exposición a peligros o impactos con barreras) mientras preserva el rendimiento de recompensa, sin requerir reentrenamiento.

• En Gridworlds, TRAM evita con éxito peligros espaciales y se adapta a condiciones de riesgo cambiantes, mientras que las líneas base conscientes del riesgo (como RaSF) permanecen uniformemente conservadoras o no logran detectar riesgos espaciales.
• En Reacher, TRAM redujo las tasas de fallo del 85% al 26% cuando se introdujo una nueva barrera después del entrenamiento, con un impacto mínimo en la distancia a la meta.
• En Safety-Gymnasium, TRAM logró la probabilidad más baja de impacto con barreras (0.209 frente a 0.502 para las líneas base) y un CVaR normalizado competitivo.
• En la alineación de LLM, TRAM produjo respuestas que fueron preferidas por comparaciones pareadas de GPT-4 sobre políticas de transferencia sin riesgo, demostrando que el método puede prevenir la sobre-optimización del modelo de recompensa en espacios de texto de alta dimensión.

5. Significado y Afirmaciones

El artículo afirma que TRAM ofrece una capa de despliegue modular para sistemas de RL. Su significado principal radica en desplazar la decisión de riesgo de la fase de entrenamiento a la fase de despliegue.

• Modularidad: Permite a los operadores especificar requisitos de seguridad (mapas de peligros, referencias conductuales) después de que los comportamientos fuente ya han sido aprendidos.
• Eficiencia: Evita el costo computacional de reentrenar o resolver nuevos MDPs restringidos en tiempo de prueba, confiando en su lugar en una puntuación de fuente rápida.
• Honestidad sobre el Alcance: Los autores son modestos en sus afirmaciones, declarando explícitamente que TRAM es un método sustituto. No afirma resolver exactamente el problema completo de control de ocupación de la política cosida. En su lugar, proporciona un límite medible sobre el desajuste entre el proxy del casco fuente y el riesgo realizado, convirtiéndolo en una heurística práctica para escenarios donde las especificaciones de riesgo son dinámicas y el reentrenamiento está prohibido.

El trabajo sugiere que para muchos escenarios de despliegue (robótica, navegación, LLMs), mantener una biblioteca diversa y neutra al riesgo de comportamientos y componerlos mediante puntuación en tiempo de prueba es una estrategia más robusta que entrenar políticas especializadas y conscientes del riesgo para cada posible restricción futura.