# TRAM: Test-Time Risk Adaptation with Mixture of Agents

arXiv: https://arxiv.org/abs/2408.08812 · HTML: https://mlsanity.com/paper/2408.08812/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Agent threats, guardrails and governance, one of 16 directions in LLM agents and tool use.
- Builds on 27 of its 41 references that are in the map.
- First posted on arXiv 16 Aug 2024.

## Placed in

- LLM agents and tool use / [Agent threats, guardrails and governance](https://mlsanity.com/agents/agent-threats-guardrails-and-governance/)

## Builds on (most cited first, 27 of 27)

- [RewardBench: Evaluating Reward Models for Language Modeling](https://mlsanity.com/paper/2403.13787/) (2024)
- [Successor Features for Transfer in Reinforcement Learning](https://mlsanity.com/paper/1606.05312/) (2016)
- [DARLA: Improving Zero-Shot Transfer in Reinforcement Learning](https://mlsanity.com/paper/1707.08475/) (2017)
- [Safety-Gymnasium: A Unified Safe Reinforcement Learning Benchmark](https://mlsanity.com/paper/2310.12567/) (2023)
- [Controlled Decoding from Language Models](https://mlsanity.com/paper/2310.17022/) (2023)
- [Risk-sensitive Reinforcement Learning](https://mlsanity.com/paper/1311.2097/) (2013)
- [Learning to be Safe: Deep RL with a Safety Critic](https://mlsanity.com/paper/2010.14603/) (2020)
- [Zero-Shot Task Generalization with Multi-Task Deep Reinforcement Learning](https://mlsanity.com/paper/1706.05064/) (2017)
- [Transfer in Deep Reinforcement Learning Using Successor Features and Generalised Policy Improvement](https://mlsanity.com/paper/1901.10964/) (2019)
- [Does Zero-Shot Reinforcement Learning Exist?](https://mlsanity.com/paper/2209.14935/) (2022)
- [Safe Reinforcement Learning via Curriculum Induction](https://mlsanity.com/paper/2006.12136/) (2020)
- [Risk-Sensitive Reinforcement Learning: Near-Optimal Risk-Sample Tradeoff in Regret](https://mlsanity.com/paper/2006.13827/) (2020)
- [Risk-Averse Trust Region Optimization for Reward-Volatility Reduction](https://mlsanity.com/paper/1912.03193/) (2019)
- [Dual RL: Unification and New Methods for Reinforcement and Imitation Learning](https://mlsanity.com/paper/2302.08560/) (2023)
- [Variance Reduction for Reinforcement Learning in Input-Driven Environments](https://mlsanity.com/paper/1807.02264/) (2018)
- [Mean-Variance Policy Iteration for Risk-Averse Reinforcement Learning](https://mlsanity.com/paper/2004.10888/) (2020)
- [Entropic Risk Measure in Policy Search](https://mlsanity.com/paper/1906.09090/) (2019)
- [Hypernetworks for Zero-shot Transfer in Reinforcement Learning](https://mlsanity.com/paper/2211.15457/) (2022)
- [Cautious Reinforcement Learning via Distributional Risk in the Dual Domain](https://mlsanity.com/paper/2002.12475/) (2020)
- [Constraint-Conditioned Policy Optimization for Versatile Safe Reinforcement Learning](https://mlsanity.com/paper/2310.03718/) (2023)
- [Learning Robust Options](https://mlsanity.com/paper/1802.03236/) (2018)
- [Risk-Aware Transfer in Reinforcement Learning using Successor Features](https://mlsanity.com/paper/2105.14127/) (2021)
- [Safe Option-Critic: Learning Safety in the Option-Critic Architecture](https://mlsanity.com/paper/1807.08060/) (2018)
- [Variance Penalized On-Policy and Off-Policy Actor-Critic](https://mlsanity.com/paper/2102.01985/) (2021)
- [Probabilistically Safe Policy Transfer](https://mlsanity.com/paper/1705.05394/) (2017)
- [SF-DQN: Provable Knowledge Transfer using Successor Feature for Deep Reinforcement Learning](https://mlsanity.com/paper/2405.15920/) (2024)
- [Situational Awareness by Risk-Conscious Skills](https://mlsanity.com/paper/1610.02847/) (2016)

## Built on by (most cited first, 1 of 1)

- [Why AI systems don't learn and what to do about it: Lessons on autonomous learning from cognitive science](https://mlsanity.com/paper/2603.15381/) (2026)

## Abstract (arXiv)

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.