TRAM: Test-Time Risk Adaptation with Mixture of Agents

Technical Summary: TRAM (Test-Time Risk Adaptation via Mixture of Agents)

1. Problem Formulation

The paper addresses a critical gap in Reinforcement Learning (RL) deployment: safety specifications often change after training. An agent may be trained on a set of source tasks but deployed in an environment with new hazard maps, revised risk thresholds, or different behavioral alignment constraints (e.g., a new reference behavior).

Existing risk-sensitive RL methods typically fix the risk model (e.g., variance, CVaR, or a specific constraint) during the training phase. This approach is brittle because:

1. Over-specialization: Training with a specific risk proxy (like return variance) can collapse the diversity of source policies, eliminating behaviors that might be safe under a different, future risk specification.
2. Inadequate Risk Representation: Variance-based risks fail to capture spatial hazards (e.g., proximity to a barrier), behavioral drift from a reference, or local volatility, especially in deterministic settings where trajectory return variance is zero.
3. Retraining Costs: Adapting to a new risk specification usually requires retraining or fine-tuning, which is computationally expensive and often infeasible once a model is deployed.

The paper formalizes the problem as zero-update deployment-time adaptation. The goal is to reuse a fixed library of risk-neutral source policies $\{\pi_j\}$ to satisfy a newly specified risk functional $\rho(\pi)$ at test time without modifying any source parameters. The objective is to maximize expected return $J(\pi)$ subject to a risk constraint $\rho(\pi) \le \delta$, approximated via a penalized surrogate:
$$\max_{\pi} \tilde{J}_c(\pi) = J(\pi) - c \rho(\pi)$$

2. Methodology: TRAM

The proposed method, TRAM, decouples source training from risk specification. It operates in two phases:

Phase 1: Risk-Neutral Source Training
Source policies are trained on related source tasks using standard, risk-neutral objectives. This preserves a diverse "behavioral hull," ensuring that the library contains a wide range of occupancy distributions (state-action visitation patterns) that can be recombined later.

Phase 2: Test-Time Source Scoring and Composition
At deployment, TRAM does not retrain. Instead, it constructs a deployed policy $\pi_T$ by selecting actions from the source library based on a source-scored composition rule.

1. Risk Evaluation: For each source policy $\pi_j$, TRAM computes a cached, global occupancy-based risk score $\rho(d_{\pi_j})$. This allows for flexible risk functionals defined over discounted state-action occupancies $d_\pi$, such as:
   • Barrier Exposure: Time spent near a hazard region.
   • Behavioral Divergence: KL-divergence from a safe reference behavior.
   • Local Volatility: Per-step fluctuations in the occupancy.
2. Value Estimation: The target-task value $Q^{\pi_j}_T(s, a)$ for each source is estimated at the current state $s$. This can be done efficiently using Successor Features (SFs) via a dot product $Q^{\pi_j}_T(s, a) = \psi_{\pi_j}(s, a)^\top w_T$, or via rollouts/critics.
3. Action Selection: The deployed action is selected by maximizing the risk-adjusted source score:
   $$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)$$
   The active source can change along the trajectory, creating a "locally stitched" policy that responds to the current state and the specific risk specification.

Theoretical Guarantees:
The paper explicitly frames TRAM as a tractable source-scored surrogate, not an exact optimizer of the full occupancy-control problem of the stitched policy.

• Localized Bounds: Theoretical analysis (Theorem 4.1) provides bounds on the suboptimality of the TRAM policy relative to the optimal risk-neutral policy. The bound depends on the localized reward mismatch between the target and the best-matching source, plus a term proportional to the risk penalty weight $c$ and the Lipschitz constant of the risk functional.
• Source-Hull Mismatch: The paper acknowledges that the realized occupancy of the stitched policy may differ from the convex hull of the source occupancies. Appendix C introduces a measurable "stitching mismatch" term ($\epsilon_{stitch}$) that bridges the gap between the cached source risks and the realized risk of the deployed policy.

3. Key Contributions

1. Formalization of Zero-Update Adaptation: The paper formalizes the setting of adapting fixed source policies to newly specified occupancy-based risks without parameter updates.
2. TRAM Algorithm: Introduction of a source-scored composition rule that supports diverse risk types (barrier, divergence, volatility) at test time, contrasting with methods tied to fixed surrogates like variance.
3. Theoretical Foundations:
   • Proof that risk-neutral source training minimizes worst-case deployment-risk regret (Theorem 4.3), preserving the widest reusable occupancy hull.
   • Localized transfer guarantees that separate reward transfer quality from the deployment-risk penalty.
   • Analysis of the "stitching mismatch" between cached source risks and realized policy risks.
4. Empirical Validation: Demonstration across four distinct domains:
   • Gridworlds: Validating expressive risk modeling (spatial hazards, per-step volatility) where variance-based methods fail.
   • MuJoCo Reacher: Reducing barrier violations in continuous control while maintaining goal performance.
   • Safety-Gymnasium: Achieving lower barrier-hit probabilities compared to frozen safety proxies and variance-based baselines.
   • LLM Alignment: Adapting fixed 7B language models to new behavioral risk constraints (KL to a safe reference) without fine-tuning, showing improved alignment over pure reward-maximization.

4. Results

Experiments show that TRAM consistently improves deployment safety (reducing hazard exposure or barrier hits) while preserving reward performance, without requiring retraining.

• In Gridworlds, TRAM successfully avoids spatial hazards and adapts to changing risk conditions, whereas risk-aware baselines (like RaSF) remain uniformly conservative or fail to detect spatial risks.
• In Reacher, TRAM reduced failure rates from 85% to 26% when a new barrier was introduced post-training, with minimal impact on the distance to the goal.
• In Safety-Gymnasium, TRAM achieved the lowest barrier-hit probability (0.209 vs. 0.502 for baselines) and competitive normalized CVaR.
• In LLM alignment, TRAM produced responses that were preferred by GPT-4 pairwise comparisons over risk-free transfer policies, demonstrating that the method can prevent reward-model over-optimization in high-dimensional text spaces.

5. Significance and Claims

The paper claims that TRAM offers a modular deployment layer for RL systems. Its primary significance lies in shifting the risk decision from the training phase to the deployment phase.

• Modularity: It allows operators to specify safety requirements (hazard maps, behavioral references) after the source behaviors are already learned.
• Efficiency: It avoids the computational cost of retraining or solving new constrained MDPs at test time, relying instead on fast source scoring.
• Honesty about Scope: The authors are modest about their claims, explicitly stating that TRAM is a surrogate method. It does not claim to solve the full occupancy-control problem of the stitched policy exactly. Instead, it provides a measurable bound on the mismatch between the source-hull proxy and the realized risk, making it a practical heuristic for scenarios where risk specifications are dynamic and retraining is prohibited.

The work suggests that for many deployment scenarios (robotics, navigation, LLMs), maintaining a diverse, risk-neutral library of behaviors and composing them via test-time scoring is a more robust strategy than training specialized, risk-aware policies for every potential future constraint.