PaperlayerView exact revision on arXiv

Source-linked AI summary

TRACE: Training Reasoning Agents for Causal Exploration with Synthesized Rewards

Rui Sun, Zhan Shi, Bing He

arXiv:2609.10315v1cs.AIcs.LG

TL;DR

Diagnostic reasoning lacks scalable objective verification because complex-data anomalies can have noisy, confounded, and ambiguous causes. TRACE engineers verification by injecting hidden interventions into a simulator, then uses the resulting oracle labels for training and evaluation. On held-out advertising diagnosis episodes, SFT followed by RL with synthesized rewards raises Qwen3.5-35B-A3B to 0.757 FullAttr@1, surpassing evaluated prompted baselines.

  • Problem

    Diagnostic reasoning over complex data lacks inexpensive objective verification because expert investigation is costly and attribution may remain uncertain amid noisy, confounded evidence.

  • Method

    TRACE injects sampled hidden interventions into a controlled digital-advertising simulator, retains oracle labels, and trains agents with synthesized rewards while they investigate through Python and SQL.

  • Results

    0.757 FullAttr@1 is reached by Qwen3.5-35B-A3B after SFT followed by RL, surpassing every evaluated prompted baseline, including Qwen3.5-122B-A10B.

  • Takeaways & Limitations

    Simulation-based verification provides a route to scalable reinforcement learning for otherwise ambiguous diagnostic reasoning tasks.

  • Takeaways & Limitations

    Further gains require more reliable multi-dimensional slice recovery and better calibration among related causes.

Abstract

from arXiv · show

Reinforcement learning with verifiable rewards (RLVR) has advanced language-model reasoning in domains such as mathematics and code, where objective answers are inexpensive to check. Diagnostic reasoning over complex data lacks this advantage: establishing the true cause of an anomaly often requires costly expert investigation and may remain ambiguous after the fact. We ask whether this asymmetry of verification can instead be engineered. We sample an intervention, inject it into a controlled simulator, and generate the observations it would produce. The hidden intervention provides an oracle label and objective reward, while the agent must still investigate noisy, confounded, and distributed evidence. We instantiate this approach in TRACE, a digital-advertising diagnostic environment with 12 root causes and fine-grained segment attribution. Agents investigate each episode using Python and SQL and must identify both the root cause and, when applicable, the affected segment assignment. On a held-out 235-episode test set, the strongest prompted baseline, Claude Opus 5, reaches 0.686 FullAttr@1. Supervised fine-tuning raises Qwen3.5-35B-A3B from 0.159 to 0.637, and subsequent RL with synthesized rewards reaches 0.757, outperforming all evaluated prompted baselines, including frontier closed-source models and a prompted Qwen3.5-122B-A10B model. The resulting policy also uses substantially fewer tool calls than the prompted 35B base. These results provide evidence that access to a scalable, objective training signal can be a more important constraint than model scale alone. More broadly, simulation-based verification can make otherwise ambiguous diagnostic reasoning tasks amenable to scalable reinforcement learning.

1 Introduction

TRACE addresses the lack of scalable, objective verification for diagnostic reasoning by engineering hidden ground-truth interventions in simulation. In digital-advertising diagnosis, SFT followed by RL with synthesized rewards substantially improves attribution and can outperform larger prompted models.

  • Motivation: Diagnostic reasoning lacks inexpensive objective verifiers because causes can be concurrent, observations noisy, and plausible explanations observationally similar.Expert investigation may remain uncertain, while learned judges can introduce reward hacking.
  • Approach: TRACE samples an intervention, injects it into a controlled simulator, and retains the hidden intervention as an oracle label for deterministic verification.The agent still investigates noisy, confounded, and distributed evidence without observing the label.
  • Approach: TRACE combines a simulator, oracle verifier, and RL pipeline in a digital-advertising environment where agents investigate campaign anomalies with Python and SQL.The environment supports root-cause attribution while ruling out realistic confounders.
  • Results: 0.757 FullAttr@1 is achieved by Qwen3.5-35B-A3B after SFT followed by RL with synthesized rewards, versus 0.159 before training and 0.686 for Claude Opus 5.The result is reported on a held-out 235-episode test set and surpasses every evaluated prompted baseline.
  • Contributions: The paper contributes simulator–oracle–RL reward synthesis, the TRACE diagnostic environment, and evidence that post-training can improve an open-weight 35B agent beyond frontier prompted models.It also studies supervised initialization and RL reward design through ablations.

2 The TRACE Environment

TRACE generates difficult digital-advertising diagnosis episodes by injecting hidden causes into noisy, confounded campaign data. Its taxonomy and controllable complexity require agents to identify what changed, where it occurred, and when it began using multi-table evidence.

  • Environment: TRACE episodes pair a campaign’s observable history with a hidden injected intervention, and the environment includes a simulator, task interface, and oracle verifier.The oracle certifies solvability and produces reward labels while the intervention defines the ground-truth cause.
  • Diagnostic challenge: Diagnosis requires resolving cause, location, and timing because metric shifts may have multiple explanations, be diluted in aggregates, or emerge gradually.These three questions define the core reasoning pattern TRACE is designed to elicit.
  • Data generation: TRACE generation has three stages: building a campaign population, sampling an intervention, and simulating the resulting metrics.The sampled intervention includes a cause type, driver slice, signal strength, and onset profile.
  • Data generation: Active causes affect generated data through total impression volume, segment allocation, and per-segment rates.The simulator uses baseline campaign properties, seasonality, intervention multipliers, and stochastic segment allocation.
  • Cause taxonomy: Segment-mix causes redistribute impressions while leaving campaign totals roughly unchanged, whereas segment-specific causes alter rates within localized slices and are hardest to identify.Segment-specific diagnosis requires localizing the driver slice and matching a multi-metric signature.
  • Cause taxonomy: No-signal episodes contain only stochastic noise or confounders and test whether agents abstain when evidence is insufficient.The benchmark also varies dimensional complexity through one- or two-attribute slices and temporal complexity through immediate, delayed, or ramping onset profiles.
  • Agent interface: Agents query campaign- and segment-level fact tables through Python backed by SQL, while ground-truth labels remain separate for verification, scoring, and reward.Segment-specific answers additionally require an affected segment slice.

3 RL with Synthesized Rewards

TRACE trains diagnostic agents with rewards synthesized by comparing final attributions against hidden oracle labels. Its reward design combines partial attribution credit, complete-attribution credit, formatting validity, and a shared tool-use setup for training and evaluation.

  • Reward synthesis: TRACE pairs each agent-visible diagnostic task with a hidden oracle label, enabling objective synthesized rewards without human annotation or an LLM judge.The agent-visible task contains the observed data, while the label records the injected intervention.
  • Training setup: Training uses GRPO over groups of sampled trajectories, with 5,000 oracle-verified episodes and the same Python-and-SQL tool interface used for evaluation.The dataset contains 4,472 training and 528 validation tasks, while evaluation uses 235 held-out tasks.
  • Reward design: The reward combines graded attribution, binary full attribution, and formatting terms, with nonnegative weights summing to one so r ∈[0, 1].The main configuration uses (wattr, wfull, wfmt) = (0.65, 0.30, 0.05).
  • Reward design: The attribution reward gives credit only when the predicted cause is correct and scores the affected segment slice when one is required.Slice agreement uses Jaccard similarity, while episodes without a required driver slice set rslice = 1.
  • Reward design: The full-attribution reward equals one only for an exactly correct cause and driver slice, complementing the dense partial credit from graded attribution.This distinguishes incomplete multi-dimensional slice recovery from complete attribution.
  • Reward computation: Reward computation parses the final root cause and, when required, affected segment slice; evidence supports interpretability but does not affect reward.This keeps training and evaluation focused on the same attribution target.

4 Experiments

TRACE evaluates whether synthesized rewards improve tool-using diagnostic attribution beyond prompting, across held-out episodes, training variants, reward designs, and tool-use efficiency. RL after supervised fine-tuning achieves the strongest overall attribution, while two-dimensional slices and no-signal calibration remain difficult.

  • Experimental setup: The held-out evaluation uses 235 episodes, emphasizes exact segment attribution, and measures cause correctness, full attribution, slice agreement, no-signal accuracy, parsing, and tool use.The benchmark includes segment-specific cases and requires exact driver-slice matches when applicable.
  • Main results: 0.686 FullAttr@1 is achieved by Claude Opus 5, the strongest prompted baseline, showing substantial remaining errors on segment-specific episodes.TRACE is described as challenging and unsaturated for prompted frontier models.
  • Main results: SFT→RL reaches 0.757 FullAttr@1, improving 12.0 percentage points over SFT’s 0.637 and surpassing all prompted baselines.The result supports an additive contribution from synthesized-reward RL after supervised fine-tuning.
  • Main results: 0.757 versus 0.283 FullAttr@1 shows the post-trained 35B model substantially outperforming the prompted Qwen3.5-122B-A10B model and every evaluated closed-source baseline.The comparison indicates that post-training can outweigh prompted model scale in this setting.
  • Training ablations: Adding the full-attribution reward raises overall FullAttr@1 from 0.596 to 0.757, with one-dimensional slices improving from 0.67 to 0.92 and two-dimensional slices from 0.00 to 0.27.The gains are especially pronounced when exact attribution requires recovering multiple driver dimensions.
  • Performance breakdown: Two-dimensional attribution remains the principal challenge: SFT→RL improves FullAttr@1 from 0.04 to 0.27, below Claude Opus 5 at 0.33, while no-signal accuracy falls from 0.76 after SFT to 0.49.The breakdown also reports recurring incomplete slices, over-attribution without signal, and confusion among related causes.
  • Tool-use efficiency: SFT reduces mean executed tool calls from 22.05 to 10.75, while RL after SFT adds only 0.98 calls to reach 11.73 as FullAttr@1 improves by 12.0 percentage points.The comparison indicates that post-training gains are not explained simply by making more tool calls.

5 Conclusion

The paper introduces simulator–oracle–RL training for diagnostic reasoning when natural verifiers are scarce, instantiating it in TRACE. On held-out evaluation, synthesized-reward RL improves attribution beyond prompted baselines, supporting simulation-based verification as a scalable training route.

  • A simulator samples hidden interventions, generates resulting data, and preserves oracle labels, making difficult diagnostic attributions objectively verifiable despite noise and confounding.
  • TRACE is a generative digital-advertising environment containing campaign-wide, segment-mix, and segment-specific root causes.
  • 0.757 FullAttr@1 is achieved after RL with synthesized rewards, compared with 0.637 after SFT and 0.686 for the strongest prompted baseline.
  • The results provide evidence that effective post-training signals, including scalable objective rewards, may constrain performance more than model scale alone in this diagnostic setting.
  • A controllable generative model can provide both the intervention that creates an ambiguous diagnostic problem and its verifier.

A Database and Output Schemas

TRACE separates agent-visible fact tables from oracle-only labels and uses structured verification to evaluate causal and segment attribution. The final answer requires a root cause, evidence, and a driver-segment assignment when the episode is segment-specific.

  • Agents query a DuckDB fact store through Python and SQL, while the tool-visible database excludes ground-truth fields.
  • The agent-visible schema consists of four fact tables with date and campaign identifiers as appropriate.
  • Oracle-only tables record the injected cause, affected segment, intervention properties, associated evidence, and verifier checks.
  • The required output contains a non-empty root_cause, optional driver_segments, an evidence list, and a textual explanation.
  • The verifier compares the episode with an equal-length preceding baseline and checks cause-specific signatures, including campaign, mix, or injected-segment changes.
  • Segment-specific signals must exceed their own baseline temporal variation with a minimum signal-to-noise ratio of 1.
  • Alternative causes are tested against the same visible data and episodes are rejected when the injected signal is absent or a matching alternative cannot be eliminated.

B.1 Training Setup

Training uses a Megatron-LM, SGLang, and Ray stack for asynchronous optimization and rollout inference. Supervised fine-tuning uses oracle-filtered teacher trajectories while preserving complete tool-use interactions.

  • Qwen3.5-35B-A3B training combines Megatron-LM optimization, SGLang rollout inference, and Ray orchestration.
  • Asynchronous RL overlaps next-step rollout generation with policy optimization and synchronizes updated weights after each optimizer step.
  • SFT trains for three epochs on 1,200 oracle-filtered teacher trajectories with a 10^-5 learning rate and cosine decay.
  • Each SFT example retains the complete interaction, including prompts, tool calls, tool outputs, and the final structured answer.

B.3 Reinforcement Learning

The reinforcement-learning setup uses shared GRPO sampling and a fixed training budget, with a separate KL ablation against the SFT reference policy. Validation is stratified across causal and attribution difficulty factors and kept disjoint from training data.

  • Each optimizer step samples 32 prompts and eight trajectories per prompt under a 32,768-token context window.
  • The agent uses a 30-turn backstop and at most 2,048 generated tokens per turn, while rollout temperature is 1.0 for within-group exploration.
  • The KL ablation penalizes divergence from the SFT reference policy with coefficient 0.005 using the low-variance k3 estimator.
  • All RL conditions share the data split, rollout group size, optimizer settings, and nominal 300-step budget.
  • The validation split is stratified by root cause, signal level, and one- versus two-dimensional attribution and is disjoint from RL prompts and SFT demonstrations.

C.1 Evaluation Protocol

The evaluation standardizes model interfaces and sampling procedures, scores both root-cause and segment attribution, and separates decision parsing from stricter schema compliance. A paired sensitivity analysis tests whether parser choice changes the reported results.

  • Evaluation setup: All models use the same task specification, candidate-cause definitions, answer contract, and Python/SQL interface over a persistent agent-visible fact store.Tool calls time out after 60 seconds, outputs truncate at 8,000 characters, and a final-turn nudge requests an answer before the interaction budget ends.
  • Sampling and budgets: Open-weight evaluations sample five trajectories per task with temperature 0.6 and a 60-step backstop, while API models use provider stochastic sampling and a 30-turn backstop.Claude and GPT baselines use xhigh reasoning effort; their output-token budgets are 32,768 and 16,384, respectively.
  • Metrics: Cause@k counts root-cause correctness, whereas FullAttr@k additionally requires an exact segment assignment for segment-specific episodes.At k = 1, the estimator is the mean single-trajectory success probability across test tasks.
  • Decision parsing: The primary parser extracts a non-empty root cause and any segment assignment, while malformed report-only fields do not invalidate an otherwise scoreable decision.A malformed segment assignment is treated as absent and cannot earn full attribution on a segment-specific episode; decision-parse rate is reported separately.
  • Sensitivity analysis: Table 7 re-scores the same trajectories with an exact-schema parser and the main decision parser to test parser sensitivity.Exact-schema parsing leaves SFT and SFT→RL results unchanged but underestimates FullAttr@1 when report-only fields are malformed or missing.
Loading 2609.10315v1…