RUBAS: Rubric-Based Reinforcement Learning for Agent Safety

2026-06-02

AI Summary

This work addresses the challenge of balancing safety and utility when deploying large language models as tool-using agents in real-world tasks. The authors propose a reinforcement learning framework grounded in a multidimensional, structured scoring rule that decomposes tool-use behavior into four interpretable dimensions: tool invocation safety, parameter safety, response safety, and usefulness. By constructing trajectory-level reward signals based on this decomposition, the framework enables fine-grained alignment with safety and utility objectives. Experimental results demonstrate that the method significantly enhances agent safety across multiple benchmarks, effectively suppresses tool-related hallucinations, and maintains task performance comparable to baseline approaches.

Abstract

The evolution of LLMs into tool-enabled agents creates a new class of safety challenges associated with real-world execution rather than simple text generation. Existing alignment methods often rely on coarse refusal signals or static supervision, making it difficult to balance safety with useful tool execution across diverse agentic risks. We introduce RUBAS, a rubric-based reinforcement learning framework for agent safety. RUBAS decomposes agent behavior into four dimensions: tool-use safety, argument safety, response safety, and helpfulness. These structured rubrics provide fine-grained and interpretable rewards over complete agent trajectories, enabling reinforcement learning to optimize safe tool use while preserving task completion. Extensive experiments across multiple agent safety benchmarks and models show that RUBAS improves safety over standard alignment baselines, reduces tool-grounded hallucinations, and maintains competitive utility. Our results suggest that multi-dimensional rubric rewards provide an effective training signal for aligning LLM agents in safety-critical tool-use settings.

Problem

• agent safety
• tool-use safety
• LLM alignment
• reinforcement learning
• safety-critical tool use

Innovation

• rubric-based reinforcement learning
• agent safety
• tool-use safety
• multi-dimensional rewards
• LLM alignment