PaperlayerView exact revision on arXiv

Source-linked AI summary

RUBAS: Rubric-Based Reinforcement Learning for Agent Safety

Xian Qi Loye, Qinglin Su, Zhexin Zhang, Shiyao Cui, Qi Zhu, Fei Mi, Hongning Wang, Minlie Huang

arXiv:2606.04051v1cs.LGcs.AIcs.CR

TL;DR

Tool-enabled agents face safety challenges across real-world execution, while existing alignment often relies on coarse rewards or static supervision. RUBAS uses four-dimensional rubrics to score complete trajectories, and experiments report improved safety with preserved utility and stronger hallucination resistance. The paper cautions that synthetic training scenarios may not reflect deployment ambiguity and that novel adversarial cases may remain challenging.

  • Problem

    Existing alignment methods struggle to balance safety and useful tool execution across agentic risks, and safety must be assessed across actions and interactions, not only textual responses.

    Source: §1 · p. 1

  • Method

    RUBAS uses structured rubric criteria across tool-use safety, argument safety, response safety, and helpfulness to provide rewards over agent trajectories.

    Source: §1 · p. 2, §7 · p. 8

  • Results

    Across benchmarks and backbone models, RUBAS improves safety while preserving utility and shows stronger hallucination resistance.

    Source: §7 · p. 8

  • Takeaways & Limitations

    RUBAS supports training agents by supervising intermediate tool-interaction decisions as well as final responses.

    Source: §7 · p. 9

    RUBAS may still fail on novel adversarial cases outside its training taxonomy, so deployment should combine it with additional safeguards.

    Source: Limitations · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Tool-enabled agents create safety risks across entire interactions, but existing alignment methods struggle to jointly balance safe actions and useful task completion. RUBAS addresses this with four-dimensional trajectory rubrics that provide structured rewards for reinforcement learning.

1 Introduction

  • Agent safety spans tool-use safety, argument integrity, response safety, and helpful completion across multi-turn interactions.The threat sources include risky or underspecified instructions, agent failures during tool use, and untrusted tool or external-system feedback.Source: §1 · p. 1
  • Existing SFT and coarse-reward RL methods struggle to jointly capture safe behavior and helpfulness across complex agent settings.SFT may imitate surface patterns that degrade across settings, while binary or sparse RL rewards miss nuanced safety–utility trade-offs.Source: §1 · p. 1
  • RUBAS models those dimensions jointly as structured sub-goals, rather than treating safe and useful behavior as a single binary decision.The four dimensions explicitly cover tool use, argument handling, response quality, and task completion.Source: §1 · pp. 1–2
  • Scorer functions turn checklist-style, verifiable criteria into a dense scalar reward for jointly optimizing safety and task completion.RUBAS uses the relative ease of verifying an output compared with generating it as a feedback mechanism.Source: §1 · p. 2
  • Experiments report reduced harmfulness and hallucinations while preserving agents’ ability to invoke tools effectively.The introduction presents this as the outcome of extensive experiments.Source: §1 · p. 2

2 Related Work

Prior work covers agent-safety evaluation, SFT and RL alignment, and rubric-based rewards, but gaps remain in everyday-risk assessment and trajectory-level safety alignment. RUBAS applies multi-scenario safety data and comprehensive rubrics to that underexplored area.

2.1 Safety Evaluation of LLM-based Agents

  • Existing benchmarks often emphasize explicitly malicious scenarios, leaving subtle risks in benign contexts and dual-use-tool misuse less assessed.The review says this limits assessment of generalization to everyday, ambiguous application scenarios.Source: §2.1 · p. 2

2.2 Safety Alignment for LLM-based Agents

  • SFT depends on intensive supervision, while coarse rule-based RL rewards can encourage overly conservative refusal and offer limited safety–utility guidance.The cited review describes these as constraints on alignment in sophisticated agent interactions.Source: §2.2 · p. 2

2.3 Reinforcement Learning with Rubric-based Rewards

  • Rubric-based rewards offer fine-grained, interpretable supervision, but their use for agent safety remains underexplored.RUBAS extends this direction with multi-scenario safety data and comprehensive trajectory-alignment rubrics.Source: §2.3 · p. 2

3 Rubric System

RUBAS evaluates complete agent trajectories across safety and helpfulness dimensions, then aggregates rubric criteria into rewards that can emphasize different behaviors and enforce critical constraints.

3.1 Rubric Design

  • RUBAS replaces a single holistic judgment with four dimensions: tool-use safety, argument safety, response safety, and helpfulness.Together, they target user-driven, tool-execution, and environmental risks.Source: §3.1 · p. 3
  • Its trajectory-aware rubric jointly evaluates intermediate tool choices, tool arguments, and the final response using binary, verifiable criteria.Reusable templates can instantiate checks such as whether a required tool was called, bound to a scenario-specific tool.Source: §3.1 · p. 3

3.2 Reward Framework

  • The rubric reward aggregates dimension scores, each a normalized weighted average of that dimension’s criteria.Criterion weights control within-dimension importance, while dimension weights are fixed within each data type and may disable a dimension when set to zero.Source: §3.2 · p. 3
  • Criterion satisfaction is mapped from [0, 1] to [−1, 1], and criterion-specific scorers assess the relevant part of the trajectory.The passage describes the mapping and scorer role as components of reward computation.Source: §3.2 · p. 3
  • Each criterion scorer checks the relevant trajectory component, with deterministic programmatic checks used except for the refusal-judgment scorer, which uses a judge model.The checked component may be tool selection, tool arguments, or the final response.Source: §3.2 · p. 3
  • Non-negotiable criteria penalize violations, while larger criterion weights make high-impact behaviors dominate their dimension scores.Examples include invoking the correct tool and avoiding unsafe arguments.Source: §3.2 · p. 3

4 Method

RUBAS builds safety tasks and context-specific rubrics, scores complete trajectories with completeness and reasoning safeguards, and uses group-relative advantages to update its policy.

4.1 Prompt and Environment Collection

  • The task collection combines 82 environments and 805 tools with risk categories and patterns to synthesize harmful, sensitive, and benign instances.GPT-5.1 generates user instructions and execution environments from the task specifications.Source: §4.1 · p. 3

4.2 Rubric Generation

  • For each instance, GPT-5.1 selects relevant rubric dimensions and generates scenario-specific, verifiable criteria from a shared template.The rubric is tailored to task context while keeping a consistent evaluation format.Source: §4.2 · pp. 3–4

4.3 Reward Function

  • Incomplete trajectories receive a fixed penalty instead of rubric scoring, and complete trajectories receive an additional penalty when the final reasoning trace is empty.Completeness requires reaching the stop condition; the empty-reasoning penalty is −0.6.Source: §4.3 · p. 4

4.4 Training Methodology

  • The policy samples multi-turn trajectories per prompt, scores them with instance-specific rubrics, and updates through GRPO using group-relative advantages.A trajectory includes reasoning, tool calls and arguments, observations, and the final response.Source: §4.4 · p. 4

5 Experiments

The experiments compare RUBAS with rule-based, judge-based, preference, and supervised baselines across agent-safety, utility, and hallucination evaluations. RUBAS achieves strong safety–utility results across models, while the analyses contrast its trajectory-level rewards and online training with coarser or static supervision.

5.1 Experimental Setting

  • Experiments evaluate Qwen3-8B, Qwen3-14B, and GLM-4.7-Flash in thinking mode on four safety benchmarks, BFCL utility across 2,501 Live and Non-Live cases, and ToolBeHonest hallucination across 700 samples.The 2,353-instance training set contains 2,000 harmful, 235 sensitive, and 118 benign cases.Source: §5.1 · pp. 4–5

5.2 Baselines

  • The four baselines are Rule’s binary indicators, GuardModel’s external binary judge, DPO’s RUBAS-scored preference pairs, and SFT on filtered safe trajectories.SFT trajectories are generated with Gemini 2.5 Flash and filtered with ShieldAgent.Source: §5.2 · p. 5

5.3 Results

  • On GLM-4.7-Flash, RUBAS improves average risk from 37.1% to 28.2% and achieves 71.0% utility, although its safety gain is smaller than on the Qwen models.Rule and GuardModel degenerate on this model; overall, RUBAS remains stable across model families.Source: §5.3 · pp. 5–6
  • Table 1 reports the benchmark comparison: on Qwen3-8B, RUBAS lowers average risk from 52.7% to 15.9%, versus 24.2% for Rule and 26.0% for GuardModel, with 79.5% utility.RUBAS has the lowest average risk among non-degenerate methods across all three models.Source: §5.3 · p. 5, Table 1 · p. 6
  • Rule and GuardModel provide coarse signals that miss trajectory details, whereas RUBAS’s multi-dimensional feedback yields better Qwen3-8B safety at comparable utility and retains a meaningful signal on GLM-4.7-Flash.The coarse baselines both degenerate on GLM-4.7-Flash.Source: §5.3 · p. 6
  • SFT raises Qwen3-8B InjecAgent ASR from 5.5% to 21.9%; DPO’s fixed base-model preference pairs can mismatch later policies, and DPO often harms utility or collapses.SFT increases indirect prompt-injection ASR on all three models.Source: §5.3 · p. 6
  • RUBAS combines interpretable rubric feedback with online GRPO updates as errors emerge; the paper attributes its best safety–utility tradeoff across three models to their complementarity.All models except the DPO and SFT baselines are trained using GRPO.Source: §5.1 · p. 5, §5.3 · p. 6
  • RUBAS achieves the highest ToolBeHonest hallucination scores on all three models; its rubric evaluates tool arguments and whether final responses are grounded in available evidence.The authors interpret this as improved safety and tool-grounded reliability, rather than simply more frequent refusal.Source: §5.3 · p. 7

6 Analysis

The analysis examines how rubric dimensions and weighting affect safety and utility, checks rubric reliability, and compares behavior in qualitative cases. Ablations show distinct effects across dimensions, while rubric variants and case studies provide additional evidence about evaluation consistency and risk handling.

6.1 Ablation Study

  • Across models, FIXED weighting matches or improves safety over GEN by +0.1 to +3.0 points, while utility differences stay within [−1.0, +0.8] points; Figure 2 compares FIXED minus GEN.FIXED’s largest safety gain is +3.0 points on GLM-4.7-Flash; the authors adopt it for simplicity and reproducibility.Source: Figure 2 · p. 7, §6.1 · p. 7
  • Removing tool-use safety reduces safety by 5.7 points to 78.4% with negligible utility change, while removing argument safety raises safety to 90.3% but lowers utility by 3.8 points to 75.7%.The authors connect argument safety to appropriate tool arguments and reliable task completion.Source: §6.1 · p. 7
  • Removing helpfulness reduces utility by 2.1 points to 77.4% while slightly improving safety, consistent with its role in limiting over-refusal on benign requests.The authors conclude that the dimensions contribute distinctly to balancing safety, utility, and grounded behavior.Source: §6.1 · p. 7
  • Removing response safety causes the largest safety decline, down 12.0 points to 72.1%, with utility nearly unchanged; Figure 3 presents the dimension-removal ablation.The authors describe response safety as the final safeguard before user-facing output.Source: §6.1 · p. 7, Figure 3 · p. 8

6.3 Case Study

  • In case studies, RUBAS asks for confirmation before potentially irreversible file deletion and corrects misinformation from unreliable tool outputs, unlike the baselines described.The authors interpret these cases as evidence of risk recognition and tool-output verification beyond superficial rule adherence.Source: §6.3 · p. 8

6.2 Validation of Rubric Reliability

  • Rubric sources show broadly consistent dimension-level scoring, but the human-revised rubric is generally stricter; GPT-5.1 has W-HMAE 0.1147 and five-generation W-IMAE 0.1203.Figure 4 compares dimension-level scores across rubric variants; the authors support GPT-5.1-generated rubrics for scalable annotation.Source: Figure 4 · p. 8, §6.2 · p. 8

7 Conclusion

RUBAS evaluates complete tool-enabled agent trajectories using four safety and helpfulness dimensions, producing fine-grained rewards for diverse risks. Across benchmarks and backbone models, the paper reports improved safety, preserved utility, and stronger hallucination resistance.

7 Conclusion

  • RUBAS evaluates complete trajectories across tool-use safety, argument safety, response safety, and helpfulness instead of reducing agent safety to one binary decision.The framework targets risks from user instructions, agent tool execution, and environmental feedback.Source: §7 · p. 8
  • Across multiple benchmarks and backbone models, RUBAS improves safety while preserving functional utility and shows stronger hallucination resistance.The conclusion also reports that multi-dimensional rubrics and online reinforcement learning matter for stable and effective alignment.Source: §7 · p. 8
  • The paper presents supervising intermediate tool- and environment-interaction decisions, not only final responses, as a practical direction for safer and more useful LLM agents.This proposed direction is framed for complex tool-use settings.Source: §7 · p. 9

Limitations

RUBAS has four stated limitations: synthetic training scenarios may not reflect deployment, rubric rewards depend on their LLM generator, experiments cover thinking-mode models, and novel adversarial cases may remain uncovered.

Limitations

  • Synthetic scenarios may not reflect deployed agents’ ambiguous or evolving intent, external states, side effects, and domain-specific policies.The authors call for further validation with live tools and authentic human instructions.Source: Limitations · p. 9
  • Reward quality depends on the rubric generator’s capability, calibration, and biases, despite encouraging agreement with human-calibrated rubrics.The reliability analysis is reported in Section 6.2.Source: Limitations · p. 9
  • Experiments focus on thinking-mode models, leaving behavior on non-thinking models and systems without separable reasoning traces unexamined.The authors note that this setting matches many recent agentic LLMs with structured intermediate reasoning.Source: Limitations · p. 9
  • RUBAS is not a complete safety solution and may fail on novel adversarial cases outside its training taxonomy.The authors recommend pairing rubric-based alignment with runtime monitoring, human confirmation for high-impact actions, continuous red-teaming, and other safeguards.Source: Limitations · p. 9

Ethical Considerations

The paper describes dual-use risk and data safeguards, and notes that the rubric is largely programmatically verifiable apart from refusal checking.

Ethical Considerations

  • Synthetic safety scenarios could expose high-risk requests or unsafe tool-use patterns, so the paper reports aggregate results without operational attack details.The remaining data is synthetically generated from public research artifacts, benchmark metadata, and code, with no private user data or identifying personal information.Source: Ethical Considerations · p. 9

A Rubric Design Details

  • The rubric criteria cover tool-use safety, argument safety, response safety, and helpfulness, with a representative template mapping dimensions to binary criteria.The table of fixed dimension weights is referenced, but the supplied passages do not provide its values.Source: A Rubric Design Details · p. 12
  • Most rubric criteria are programmatically verifiable; must_refuse is the sole criterion not fully verifiable this way.The design therefore remains largely verifiable, with the exception isolated to refusal checking.Source: A Rubric Design Details · pp. 12–13

B Details of Data Construction

Data construction spans harmful, sensitive, and benign tasks, pairs risk domains with execution patterns, and uses generated tasks and rubrics grounded in a broad synthesized tool collection.

B Details of Data Construction

  • The taxonomy separates 10 harm categories from six execution failure patterns, covering explicit malicious intent and subtler agent-reasoning failures.Risk patterns identify where failures arise in the agent workflow, such as embedding untrusted input in a tool argument.Source: B Details of Data Construction · pp. 13–14
  • Harmful and sensitive instances sample risk categories and patterns with relevant tools, while benign instances use neutral tool subsets to preserve helpfulness.Some harmful or sensitive tasks leave the risk category unspecified so the generator can produce emergent harms.Source: B Details of Data Construction · pp. 13–14
  • The generated environment collection comprises 82 environments and 805 unique tools, with harmful, sensitive, and benign tasks grounded in selected tool subsets.Task instructions and execution environments are synthesized from the sampled specifications.Source: B Details of Data Construction · p. 14
  • GPT-5.1 generates a task-specific rubric for each instance, using four dimensions and verifiable criteria tailored to the task context.Post-generation checks discard criteria referring to nonexistent tool parameters, grounding reward signals in the environment.Source: B Details of Data Construction · p. 14

C Details of Tool Collection

  • Tool collection uses lightweight environment-class and tool-name metadata from four datasets, then synthesizes new tool descriptions and Python implementations.Benchmark prompts, implementations, states, trajectories, and labels are not imported into training data.Source: C Details of Tool Collection · pp. 14–15
  • Training excludes held-out Agent-SafetyBench and InjecAgent test prompts and limits overlap by using name-level metadata rather than benchmark content.For Agent-SafetyBench, RUBAS uses 565 of 2,013 tool names; one implementation is used only as a one-shot generation-format example.Source: C Details of Tool Collection · p. 15

D Generation, Training, and Judging Details

The paper specifies generation, optimization, preference-data, judge, and evaluation procedures for RUBAS and its baselines.

D Generation, Training, and Judging Details

  • GPT-5.1 generates instructions and rubric criteria at temperature 0.6 with a 4,096-token limit, using separate templates for benign, harmful, and sensitive tasks.The prompt templates are referenced as Figures 6–8.Source: D Generation, Training, and Judging… · p. 16, Figure 7 · p. 17, Figure 8 · p. 17
  • GRPO training uses eight responses per prompt, batch size 32, temperature 1.0, up to four turns, and three epochs, with learning rate 8 × 10^-6.DPO and SFT use the same learning rate, batch size, turn limit, per-turn token limit, and epoch count for comparison.Source: D Generation, Training, and Judging… · p. 16
  • DPO forms preference pairs by scoring 12 Qwen3-8B responses per instance with RUBAS and selecting the highest- and lowest-scoring candidates.Candidate sampling uses temperature 1.0.Source: D Generation, Training, and Judging… · p. 16
  • SFT trajectories are generated with Gemini 2.5 Flash and filtered by ShieldAgent to retain safe, high-quality examples.Source: D Generation, Training, and Judging… · p. 16
  • The GuardModel baseline uses Qwen3Guard-Gen-8B’s safety judgment on the instruction and final response, while RUBAS uses it only for must_refuse.The Rule setting and DPO data filtering also use the same refusal judgment.Source: D Generation, Training, and Judging… · p. 18
  • Models are labeled degenerate when BFCL is below 0.5% and most cases show repetitive text or no tool calls; initially degenerate runs are repeated with two additional seeds.A degenerate label requires all three seeds to remain degenerate.Source: D Generation, Training, and Judging… · p. 18
  • Training uses eight NVIDIA A800 GPUs; a Qwen3-8B GRPO run takes about 3.5 hours, versus about one hour for DPO or SFT.GRPO uses AReaL, while SFT and DPO use TRL.Source: D Generation, Training, and Judging… · p. 18

E Details of Rubric Reliability Validation

The appendix validates rubric reliability through cross-model agreement, human alignment, and repeated-generation consistency. Generated rubrics preserve broad safety-sensitive scoring patterns, while human calibration and repeated-run analyses assess their alignment and stability.

E Details of Rubric Reliability Validation

  • Each instruction was scored using the average of eight independently scored policy rollouts, then category scores were averaged across instructions.The rollout procedure compares rubric variants under the same instructions and sampled rollouts.Source: E Details of Rubric Reliability… · p. 18
  • Across rubric sources, dimension-level scoring patterns were broadly consistent, especially for tool-use safety and response safety.The comparison covers the human-revised reference and model-generated variants.Source: E Details of Rubric Reliability… · p. 18
  • Human revision primarily tightened overly broad, under-specified, or insufficiently safety-sensitive rubric criteria.The human-revised rubric generally scored more strictly than generated variants.Source: E Details of Rubric Reliability… · pp. 18–19
  • Figure 9 shows the same category ordering across rubric variants: Benign and Sensitive examples receive positive rewards, while Harmful examples receive negative rewards.Model-generated rubrics assign higher scores in most cases than the human-revised reference.Source: E Details of Rubric Reliability… · p. 18, Figure 9 · p. 19
  • Human alignment uses category-level mean absolute error against the human-revised reference, aggregated with fixed weights of 0.25, 0.25, and 0.50 for Benign, Sensitive, and Harmful examples.The validation sampled 200 instances: 50 Benign, 50 Sensitive, and 100 Harmful.Source: E Details of Rubric Reliability… · pp. 18–19, Table 6 · p. 19
  • GPT-5.1 had the smallest distance to the human-revised reference, with W-HMAE = 0.1147; Qwen3.5-397B-A17B was closest on Benign examples but deviated more on Sensitive and Harmful examples.The results support using GPT-5.1-generated and filtered rubrics for training-data construction.Source: Table 6 · p. 19, E Details of Rubric Reliability… · p. 19
  • Across five repeated GPT-5.1 rubric-generation runs, overall pairwise batch distance was W-IMAE = 0.1203, with the lowest distance on Harmful examples.Benign and Sensitive examples showed moderately larger differences.Source: Table 8 · p. 20, E Details of Rubric Reliability… · p. 20
Loading 2606.04051v1…