English

TRAM: Test-Time Risk Adaptation with Mixture of Agents

Machine Learning 2026-05-21 v2

Abstract

Deployed reinforcement learning agents often face safety requirements that are specified only after training, such as new hazard maps, revised risk thresholds, or behavioral alignment constraints. We study zero-update deployment-time adaptation, where a fixed library of risk-neutral source policies is reused under a newly specified reward-risk tradeoff. We propose TRAM (Test-Time Risk Adaptation via Mixture of Agents), a source-scored composition rule that evaluates each source policy under the target reward and an occupancy-based deployment risk, then selects actions using risk-adjusted source scores. Unlike training-time risk-sensitive methods tied to a fixed surrogate such as return variance, TRAM supports spatial barrier exposure, divergence from a reference behavior, and local volatility risks specified at test time. We explicitly characterize TRAM as a surrogate method: it does not solve the full occupancy-control problem of the stitched policy, but admits a measurable source-hull mismatch term connecting source-scored risk to realized risk. Experiments in gridworlds, MuJoCo Reacher, Safety-Gymnasium, and an LLM alignment setting show that TRAM reduces deployment risk while preserving reward, without requiring any parameter updates at test time.

Keywords

Cite

@article{arxiv.2408.08812,
  title  = {TRAM: Test-Time Risk Adaptation with Mixture of Agents},
  author = {Mohamad Fares El Hajj Chehade and Amrit Singh Bedi and Amy Zhang and Hao Zhu},
  journal= {arXiv preprint arXiv:2408.08812},
  year   = {2026}
}