The Alignment Waltz: Jointly Training Agents to Collaborate for Safety

🕒 Published (v1): 2026-01-01 · Source: ICLR · Venue: ICLR 2026 · link

Why this paper was selected

ICLR 2026; jointly training agents to collaborate, balancing helpfulness and safety

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

WALTZ RL is a multi-agent reinforcement learning framework that jointly trains a conversation agent and a feedback agent to collaboratively improve LLM safety alignment. Rather than blocking unsafe outputs with a hard safeguard, the feedback agent diagnoses problems and steers the conversation agent toward safer, less over-refusing responses. On Llama-3.1-8B-Instruct, it reduces adversarial attack success from 39.0% to 4.6% and overrefusal from 45.3% to 9.9% without degrading general capabilities.

Problem

Existing safety approaches rely on standalone safeguard models (e.g., Llama Guard) that hard-refuse any response containing unsafe content. This binary blocking (a) exacerbates overrefusal on benign but sensitive prompts, (b) discards all safe/helpful content in a partially unsafe response, and (c) provides no corrective signal. Single-model RL reduces overrefusal but remains vulnerable to adversarial attacks. No prior method simultaneously advances both the safety and helpfulness axes of the Pareto front.

Method

WALTZ RL formulates safety alignment as a cooperative, positive-sum multi-agent game between: - Conversation agent (\(\pi_c\)): generates initial and revised responses. - Feedback agent (\(\pi_f\)): reasons about the conversation agent's output, predicts unsafe/overrefuse labels, and produces a natural-language feedback string in JSON format.

Collaboration protocol: Given prompt \(p\), the conversation agent produces \(c_0\); the feedback agent produces \((f_0)\); if \(f_0\) flags a problem, the conversation agent revises to \(c_1\); repeat up to \(T_{\max}=1\) round.

Reward design: - Conversation agent reward: \(R_c = \mathbf{1}[\neg\text{unsafe} \wedge \neg\text{overrefuse}]\) (binary). - Feedback agent reward combines three terms: $\(R_f = \alpha R_f^{\text{DIR}} \cdot R_f^{\text{label}} + \lambda R_f^{\text{label}} + \gamma R_f^{\text{format}}\)$ - Dynamic Improvement Reward (DIR): \(R_f^{\text{DIR}} = R_c(c_{t+1}) - R_c(c_t)\); rewards feedback that improves the conversation agent's next revision, conditioned on label accuracy.

Two-stage training: 1. Stage 1: conversation agent frozen; feedback agent trains on all reward terms to learn format and label accuracy. 2. Stage 2: both agents co-trained; label reward \(\lambda\) set to 0 to prevent overfitting to imbalanced data (most responses are already safe late in training).

Policy gradient: Extends REINFORCE++ to two actors; multi-agent rollouts are reduced to single-actor trajectories per agent. Conversation agent training blends trajectory type-A (initial response \(c_0\)) and type-B (final revision \(c_T\) after feedback).

Inference: Both agents deploy jointly; feedback agent triggers only when it predicts unsafe=True or overrefuse=True, preserving latency on safe queries.

Key Contributions

• WALTZ RL: first framework to jointly train conversation and feedback agents end-to-end for safety alignment, deploying both at inference.
• Dynamic Improvement Reward (DIR): a temporally evolving reward for the feedback agent tied to the conversation agent's policy improvement, enabling genuine co-evolution.
• Two-stage training pipeline enabling adaptive feedback (feedback triggers only when needed).
• Empirical evidence that detailed natural-language feedback outperforms oracle template feedback, especially for correcting overrefusals.

Results

• WildJailbreak (WJ) ASR: 39.0% (baseline) → 4.6% (WALTZ RL); vs. single-model RL (13.2%), safeguard+RL (7.3%), inference-time collab (19.4%), oracle feedback (10.6%).
• OR-Bench ORR: 45.3% → 9.9%; vs. baseline+safeguard (48.7%), single-model RL (11.9%), oracle feedback (28.2%).
• Average ASR across WJ/FH/SR: 25.7% (baseline) → 3.7% (WALTZ RL).
• Average ORR across OB/FB: 25.7% → 5.4% (WALTZ RL); all baselines ≥7.6%.
• General capability (Table 2): AlpacaEval LCWR drops marginally 37.2→35.9; MMLU 68.0→68.1; TruthfulQA 37.0→37.0; IFEval metrics stable — no significant degradation.
• Feedback Trigger Rate (FTR): WALTZ RL reduces FTR on AlpacaEval to 6.7% vs. 42.6% for inference-time collaboration (no training), recovering helpfulness latency.
• Label accuracy: feedback agent label accuracy improves from 31.4% (inference-time collab) to 70.1% (WALTZ RL) on WJ.
• Adding safeguards to single-model RL worsens ORR by 6.3 pp vs. 4.1 pp when added to the untuned baseline, confirming safeguards disproportionately harm already-low-overrefusal systems.

Limitations

• Maximum feedback rounds \(T_{\max}=1\) tested in experiments; multi-round extensions increase inference latency and are not empirically explored.
• Both agents initialized from the same base model (Llama-3.1-8B-Instruct); scalability to larger or asymmetric model pairs is unvalidated.
• Training data covers only adversarial jailbreak and borderline over-refusal prompts; transfer to other safety failure modes (e.g., bias, misinformation) is not demonstrated.
• The feedback agent adds inference overhead even at \(T_{\max}=1\); FTR remains non-zero on safe queries (6.7%), introducing latency for a fraction of benign requests.
• Reward uses an LLM judge for alignment labels, inheriting whatever errors or biases that judge carries.

Relevance to Agentic AI / LLM Agents

WALTZ RL is directly relevant to multi-agent system design: it demonstrates that a critic/reviewer agent can be co-trained alongside a generator agent via RL, with the critic's reward derived from the delta in the generator's performance — a general pattern applicable beyond safety to any agent-plus-verifier architecture. The Dynamic Improvement Reward is a reusable primitive for training helper agents whose value is defined entirely by downstream agent improvement. The adaptive, inference-time deployment of the feedback agent (trigger-on-need) models a practical pattern for LLM agent pipelines where auxiliary agents must be invoked selectively to preserve throughput. For researchers building agentic systems with safety constraints, this work shows that collaborative agent training can outperform both hard safeguards and single-model RLHF on the helpfulness-harmlessness Pareto front.