Reinforcement learning finds causes of problems in simulated ad data

TRACE: Training Reasoning Agents for Causal Exploration with Synthesized Rewards

Artificial IntelligenceMachine Learning

Summary

Finding the real cause of strange things happening in complex data is often hard and expensive because experts have to check. The authors create a computer simulation where they secretly add a cause and then generate what would be observed from it. This way, an AI agent can learn to find the hidden cause by practicing in the simulation with clear right or wrong feedback. Their system called TRACE works in digital advertising data and shows it can get better at diagnosing causes than large language models that just try to guess without training. This approach shows how using simulated rewards makes hard detective work easier to teach computers.

What this means in practice

  • For digital advertising analysts: Automatically diagnose root causes of campaign anomalies using AI that learns from simulated ad data with known causes.
  • For data engineers: Build AI agents that investigate complex, noisy datasets by training them on simulated scenarios with objective feedback.

Authors

Rui Sun, Zhan Shi, Bing He

Abstract

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.