Research map / Label-efficient and robust learning / Test-time adaptation

Test-time Risk Adaptation with Mixture of Agents

NeurIPS 2026 · Label-efficient and robust learning › Test-time adaptation · approach: Realistic and stable test-time adaptation

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes TRAM, a zero-update deployment-time method that composes a fixed library of risk-neutral source policies by scoring each source action with target value minus a test-time occupancy risk penalty.

Problem

Deployed RL agents often need to satisfy safety or alignment requirements specified only after training, but retraining or constrained re-solving per specification is expensive or disallowed; fixed training-time risk models are brittle to new risk types.

Approach

Keep source policies risk-neutral, cache each source's discounted occupancy and risk once, then at each state pick the action maximizing target value minus c times the source's cached risk across sources (source-scored composition), supporting barrier, KL-divergence, and local-volatility occupancy risks specified at test time.

Key results

• In Safety-Gymnasium SafetyPointGoal0-v0, TRAM reduces barrier-hit probability from 0.502 to 0.209.
• In Reacher, TRAM cuts the failure rate from 85% to 26% relative to risk-neutral transfer while keeping distance-to-goal comparable.
• On Berkeley Nectar, the risk-free transfer policy has the highest reward-model score, but GPT-4 pairwise preferences favor TRAM overall.

Surprising findings

• The risk-free transfer policy maximizes reward-model score yet is preferred less often than TRAM by GPT-4, indicating reward over-optimization.

Limitations

• Performance depends on source-library coverage and risk-estimation fidelity, especially if all source policies are poor.
• Not an exact optimizer of the stitched policy's occupancy objective; a stitching mismatch term can be large.
• LLM experiment uses response-level risk rather than occupancy-based risk, so theory does not directly apply.

Experimental setup

Datasets: Berkeley Nectar Safety-Gymnasium SafetyPointGoal0-v0 MuJoCo Reacher

Benchmarks: Safety-Gymnasium SafetyPointGoal0-v0 MuJoCo Reacher Berkeley Nectar gridworlds

Metrics: barrier-hit probability normalized CVaR failure rate distance to goal target reward-model score GPT-4 pairwise preference

Baselines: GPI SF transfer RaSF SafeDSR-style proxy SFT-CoP-style dual baseline

Builds on: successor features generalized policy improvement risk-aware successor features

Closest papers by meaning

• Risk-Conditioned Reinforcement Learning: A Generalized Approach for Adapting to Varying Risk Measures (AAAI 2024 · similarity 0.88)
• Efficient Test-Time Adaptation For Robot Policies (NeurIPS 2026 · similarity 0.87)
• Safety Representations for Safer Policy Learning (ICLR 2025 · similarity 0.87)
• Enhancing Efficiency of Safe Reinforcement Learning via Sample Manipulation (NeurIPS 2024 · similarity 0.86)
• Safeguarded Progress in Reinforcement Learning: Safe Bayesian Exploration for Control Policy Synthesis (AAAI 2024 · similarity 0.86)
• Constraint-Adaptive Policy Switching for Offline Safe Reinforcement Learning (AAAI 2025 · similarity 0.86)