TRAM: Test-Time Risk Adaptation with Mixture of Agents

Resumo Técnico: TRAM (Adaptação de Risco em Tempo de Teste via Mistura de Agentes)

1. Formulação do Problema

O artigo aborda uma lacuna crítica na implantação de Aprendizado por Reforço (RL): as especificações de segurança frequentemente mudam após o treinamento. Um agente pode ser treinado em um conjunto de tarefas-fonte, mas implantado em um ambiente com novos mapas de perigo, limiares de risco revisados ou restrições diferentes de alinhamento comportamental (por exemplo, um novo comportamento de referência).

Os métodos existentes de RL sensíveis ao risco normalmente fixam o modelo de risco (por exemplo, variância, CVaR ou uma restrição específica) durante a fase de treinamento. Essa abordagem é frágil porque:

1. Especialização Excessiva: O treinamento com um proxy de risco específico (como variância de retorno) pode colapsar a diversidade das políticas-fonte, eliminando comportamentos que poderiam ser seguros sob uma especificação de risco diferente e futura.
2. Representação de Risco Inadequada: Riscos baseados em variância falham em capturar perigos espaciais (por exemplo, proximidade a uma barreira), desvio comportamental de uma referência ou volatilidade local, especialmente em configurações determinísticas onde a variância do retorno da trajetória é zero.
3. Custos de Retreinamento: Adaptar-se a uma nova especificação de risco geralmente exige retreinamento ou ajuste fino, o que é computacionalmente caro e frequentemente inviável uma vez que o modelo é implantado.

O artigo formaliza o problema como adaptação em tempo de implantação sem atualização. O objetivo é reutilizar uma biblioteca fixa de políticas-fonte neutras ao risco $\{\pi_j\}$ para satisfazer um funcional de risco recém-especificado $\rho(\pi)$ no momento do teste, sem modificar quaisquer parâmetros-fonte. O objetivo é maximizar o retorno esperado $J(\pi)$ sujeito a uma restrição de risco $\rho(\pi) \le \delta$, aproximado via um substituto penalizado:
$$\max_{\pi} \tilde{J}_c(\pi) = J(\pi) - c \rho(\pi)$$

2. Metodologia: TRAM

O método proposto, TRAM, desacopla o treinamento-fonte da especificação de risco. Ele opera em duas fases:

Fase 1: Treinamento Fonte Neutro ao Risco
As políticas-fonte são treinadas em tarefas-fonte relacionadas usando objetivos padrão e neutros ao risco. Isso preserva um "casco comportamental" diversificado, garantindo que a biblioteca contenha uma ampla gama de distribuições de ocupação (padrões de visitação estado-ação) que podem ser recombinadas posteriormente.

Fase 2: Pontuação e Composição de Fontes em Tempo de Teste
Na implantação, o TRAM não retreina. Em vez disso, ele constrói uma política implantada $\pi_T$ selecionando ações da biblioteca-fonte com base em uma regra de composição pontuada por fonte.

1. Avaliação de Risco: Para cada política-fonte $\pi_j$, o TRAM calcula uma pontuação de risco global baseada em ocupação em cache $\rho(d_{\pi_j})$. Isso permite funcionais de risco flexíveis definidos sobre ocupações estado-ação com desconto $d_\pi$, tais como:
   • Exposição a Barreiras: Tempo gasto próximo a uma região de perigo.
   • Divergência Comportamental: Divergência KL de um comportamento de referência seguro.
   • Volatilidade Local: Flutuações por passo na ocupação.
2. Estimativa de Valor: O valor da tarefa-alvo $Q^{\pi_j}_T(s, a)$ para cada fonte é estimado no estado atual $s$. Isso pode ser feito eficientemente usando Características Sucessoras (SFs) via um produto escalar $Q^{\pi_j}_T(s, a) = \psi_{\pi_j}(s, a)^\top w_T$, ou via simulações/críticos.
3. Seleção de Ação: A ação implantada é selecionada maximizando a pontuação de fonte ajustada ao risco:
   $$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)$$
   A fonte ativa pode mudar ao longo da trajetória, criando uma política "costurada localmente" que responde ao estado atual e à especificação de risco específica.

Garantias Teóricas:
O artigo enquadra explicitamente o TRAM como um substituto pontuado por fonte tratável, e não como um otimizador exato do problema completo de controle de ocupação da política costurada.

• Limites Localizados: A análise teórica (Teorema 4.1) fornece limites para a suboptimalidade da política TRAM em relação à política ótima neutra ao risco. O limite depende do desajuste de recompensa localizado entre o alvo e a fonte de melhor correspondência, mais um termo proporcional ao peso da penalidade de risco $c$ e à constante de Lipschitz do funcional de risco.
• Desajuste do Casco-Fonte: O artigo reconhece que a ocupação realizada da política costurada pode diferir do casco convexo das ocupações-fonte. O Apêndice C introduz um termo mensurável de "desajuste de costura" ($\epsilon_{stitch}$) que preenche a lacuna entre os riscos-fonte em cache e o risco realizado da política implantada.

3. Contribuições Principais

1. Formalização da Adaptação sem Atualização: O artigo formaliza o cenário de adaptação de políticas-fonte fixas a riscos baseados em ocupação recém-especificados sem atualizações de parâmetros.
2. Algoritmo TRAM: Introdução de uma regra de composição pontuada por fonte que suporta diversos tipos de risco (barreira, divergência, volatilidade) no momento do teste, contrastando com métodos ligados a substitutos fixos como variância.
3. Fundamentos Teóricos:
   • Prova de que o treinamento-fonte neutro ao risco minimiza o arrependimento de risco de implantação no pior caso (Teorema 4.3), preservando o casco de ocupação reutilizável mais amplo.
   • Garantias de transferência localizadas que separam a qualidade da transferência de recompensa da penalidade de risco de implantação.
   • Análise do "desajuste de costura" entre riscos-fonte em cache e riscos de política realizados.
4. Validação Empírica: Demonstração em quatro domínios distintos:
   • Gridworlds: Validação de modelagem de risco expressiva (perigos espaciais, volatilidade por passo) onde métodos baseados em variância falham.
   • MuJoCo Reacher: Redução de violações de barreira em controle contínuo enquanto mantém o desempenho de objetivo.
   • Safety-Gymnasium: Alcançando probabilidades de impacto em barreira mais baixas em comparação com proxies de segurança congelados e baselines baseadas em variância.
   • Alinhamento de LLM: Adaptação de modelos de linguagem fixos de 7B a novas restrições de risco comportamental (KL para uma referência segura) sem ajuste fino, mostrando melhor alinhamento sobre a maximização pura de recompensa.

4. Resultados

Os experimentos mostram que o TRAM melhora consistentemente a segurança de implantação (reduzindo a exposição a perigos ou impactos em barreiras) enquanto preserva o desempenho de recompensa, sem exigir retreinamento.

• Em Gridworlds, o TRAM evita com sucesso perigos espaciais e adapta-se a condições de risco em mudança, enquanto baselines conscientes de risco (como RaSF) permanecem uniformemente conservadoras ou falham em detectar riscos espaciais.
• Em Reacher, o TRAM reduziu as taxas de falha de 85% para 26% quando uma nova barreira foi introduzida após o treinamento, com impacto mínimo na distância até o objetivo.
• Em Safety-Gymnasium, o TRAM alcançou a menor probabilidade de impacto em barreira (0,209 vs. 0,502 para baselines) e CVaR normalizado competitivo.
• Em alinhamento de LLM, o TRAM produziu respostas que foram preferidas em comparações pareadas pelo GPT-4 sobre políticas de transferência sem risco, demonstrando que o método pode prevenir a superotimização do modelo de recompensa em espaços de texto de alta dimensão.

5. Significado e Alegações

O artigo afirma que o TRAM oferece uma camada de implantação modular para sistemas de RL. Seu significado principal reside em deslocar a decisão de risco da fase de treinamento para a fase de implantação.

• Modularidade: Permite que operadores especifiquem requisitos de segurança (mapas de perigo, referências comportamentais) após os comportamentos-fonte já terem sido aprendidos.
• Eficiência: Evita o custo computacional de retreinamento ou resolução de novos MDPs restritos no momento do teste, confiando em vez disso na pontuação rápida de fontes.
• Honestidade sobre o Escopo: Os autores são modestos em suas alegações, afirmando explicitamente que o TRAM é um método substituto. Não alega resolver exatamente o problema completo de controle de ocupação da política costurada. Em vez disso, fornece um limite mensurável para o desajuste entre o proxy do casco-fonte e o risco realizado, tornando-o uma heurística prática para cenários onde as especificações de risco são dinâmicas e o retreinamento é proibido.

O trabalho sugere que, para muitos cenários de implantação (robótica, navegação, LLMs), manter uma biblioteca diversificada e neutra ao risco de comportamentos e compostá-los via pontuação em tempo de teste é uma estratégia mais robusta do que treinar políticas especializadas e conscientes de risco para cada possível restrição futura.