AdaPRA: Process Reward Agents with Adaptive Rubrics for Data Science

AI summary

Overview

Research area: AI agents for data science, specifically process supervision of long-running coding agents (process reward models and agents, automatic rubric discovery, and agent harness design).

Technical level: Advanced. The paper assumes familiarity with LLM agents, tool use, code execution, and reward modeling, though the core ideas are explained in accessible terms.

Scope: The paper introduces AdaPRA, a system that pairs a persistent agent harness with a rubric-driven process reward agent whose criteria are mined from past runs and adapted per task, and evaluates it on two Kaggle-style data science benchmarks.

What This Paper Is About

Data science agents can build end-to-end machine learning pipelines, but they still trail humans on leaderboards largely because of process errors: stopping at the first valid pipeline, unsound validation, coding bugs, and wasting a fixed budget by repeating the same exploratory analysis and baselines on every attempt. Many of these errors are silent—the pipeline produces a valid submission and the trajectory shows no sign of trouble. AdaPRA addresses this with a harness that carries the workspace forward across repeated attempts on the same task, plus a process reward agent that audits both the trajectory and the workspace by executing checks, using rubrics that are discovered offline from past runs and adapted to each task.

Key Contributions

1. A persistent harness for data science agents. Unlike the DSGym harness, which resets the workspace between attempts, AdaPRA keeps one workspace per task across all attempts, so exploratory notes, trained models, validation results, and submissions persist, and later attempts receive the agent's own summaries of earlier attempts plus the final review feedback.
2. A process reward agent (PRA) that executes checks. The PRA has its own terminal and scratch space, can write and run code against the base agent's read-only files (for example, recomputing a cross-validation score from out-of-fold predictions), pulls trajectory context on demand from files rather than stuffing it into its context window, and reviews at a fixed cadence during an attempt plus a finish gate.
3. A rubric discovery procedure that mines trajectories and workspace artifacts. Two stages: an LLM judge extracts good behaviors and issues from past harness-only runs with quoted evidence and knowledge of whether each attempt ranked in the top or bottom quarter on its task; a second call clusters them into a mutually exclusive and collectively exhaustive (MECE) rubric set. Stage one mined 906 behaviors (564 issues, 342 good behaviors) from DSPredict-Easy runs, yielding 17 rubrics plus 2 cross-attempt rubrics, for 19 seed rubrics.
4. Adaptive rubrics and evidence of transfer. The PRA can add, update, and retire rubrics in a per-task store initialized with the seed set, and the paper shows that rubrics learned on mostly tabular DSPredict-Easy transfer to DSPredict-Hard's vision, text, audio, time-series, and molecular tasks, and that rubrics mined from one base model's runs improve another (DeepSeek-V3.2).

Main Findings

• Large percentile gains. Replacing the DSGym harness with AdaPRA (using Claude Sonnet-4.6 for the PRA) raises the 80B Qwen3-Coder-Next agent's mean Kaggle percentile from 44.5 to 69.0 on DSPredict-Easy and from 21.3 to 36.1 on DSPredict-Hard—24.5 and 14.8 points respectively. On DSPredict-Easy the above-median rate rises from 50.0% to 86.8% and the valid rate from 76.3% to 100.0%; on DSPredict-Hard, from 11.1% to 33.3% and from 68.5% to 90.7%.
• Supervision helps typical attempts, not just the best one. Per-attempt mean percentile rises from 36.3 to 61.0 on DSPredict-Easy and from 25.2 to 32.1 on DSPredict-Hard. A single AdaPRA attempt on DSPredict-Easy averages 61.0 percentile, versus 60.4 for a single GPT-5.1 (high) attempt in the DSGym harness.
• The harness and supervision contribute differently by difficulty. The harness alone lifts mean percentile from 44.5 to 60.8 (Easy) and 21.3 to 25.5 (Hard). The harness contributes 16.3 points and supervision 8.2 on DSPredict-Easy, versus 4.2 and 10.6 on DSPredict-Hard—i.e., easier tasks mostly need the persistent workspace, harder ones need supervision.
• Beating a competitive PRM. DataPRM (a 4B generative process reward model that scores each step 1, 0.5, or 0) reaches 54.6 best-attempt percentile on DSPredict-Easy and 23.1 on DSPredict-Hard, below AdaPRA's 69.0 and 36.1.
• Both ablated components matter. Removing code execution (the PRA reads but cannot run code) drops Easy from 69.0 to 63.0 (p=0.012) and Hard from 36.1 to 31.8 (p=0.077, 29 wins to 18 losses). Freezing the seed rubrics drops Easy by 1.7 points and Hard by 4.3 points (p=0.024; 32 wins, 13 losses).
• Rubrics transfer from easy to hard tasks. With frozen seed rubrics, supervision adds 6.3 percentage points over the harness alone on DSPredict-Hard (p=0.007), close to its 6.5-point gain on DSPredict-Easy (p=0.017). Every seed rubric fails in at least 24% of DSPredict-Hard attempts, and ranking rubrics by importance or failure rate gives nearly the same order on both benchmarks (Spearman ρ=0.91 for importance, ρ=0.66 with p=0.002 for failure rate).
• Failures are agent properties, not domain properties. The two simplicity-bias rubrics (too little model diversity and ensembling; repeating near-identical experiments without pivoting) are the most frequent failures on both benchmarks, at 68–84% of attempts. Invalid cross-validation rises from 32% to 57% of attempts on harder tasks, and unverified submissions from 33% to 52%.
• Adaptation supplies domain knowledge. The PRA adds 0.8 rubrics per task on DSPredict-Easy and 1.2 on DSPredict-Hard, adding at least one on 79% of Hard tasks versus 51% of Easy tasks. Across four adaptive runs it made 177 additions, one update, and no retirements, mostly in the first two attempts. Additions are often task-specific, such as "prediction underdispersion" on House Prices, non-negative time-to-failure predictions on LANL Earthquake Prediction, and excluding quantum-mechanical sub-components of the target in CHAMPS Scalar Coupling.
• Code execution changes what gets written. Prediction-sanity rubrics make up 35% (Easy) and 19% (Hard) of additions with execution, versus 3% and 7% without it.
• Cadence matters. With one attempt per task on DSPredict-Easy, every cadence beats the harness alone (48.4 mean percentile): every 5 turns gives 52.7, every 10 gives 56.4 (best), every 20 gives 49.9.
• It works on a second base model. With DeepSeek-V3.2 as the base agent and one attempt per task, AdaPRA raises mean percentile from 36.9 to 44.7 and the above-median rate from 23.7% to 47.4%, winning on 25 tasks and losing on 11 (p=0.005).

Methodology in Plain English

The base agent—Qwen3-Coder-Next, an 80B open-weights coding model, built on the OpenHands SDK with a terminal and file editor—works in a per-task container. Every condition, including baselines, gets five attempts per task, and Kaggle scores are used only for evaluation: neither the agent nor the PRA ever sees leaderboard scores, test labels, or other external feedback between attempts.

Instead of starting fresh, each attempt inherits the workspace from the previous one—notes, features, trained models, validation results, submissions—while only the model context resets. Later attempts also see the agent's own summaries of earlier attempts, followed by the final review feedback when supervised.

The PRA (Claude Sonnet 4.6) scores the partial trajectory and the current workspace against each rubric, producing a probability of satisfaction, an importance weight, a rationale, and a suggested next action. To keep feedback short, only relevant failing rubrics are shown: those with importance at least 0.5 and a satisfaction probability below 0.6. The PRA reviews every 10 agent turns on DSPredict-Easy and every 20 on DSPredict-Hard (to bound cost, since the PRA accounts for 75% of a Hard run's cost and Hard attempts take 1.6 times as many agent turns), plus a finish review that gates submission and allows one refinement round. Long trajectories are written as turn-annotated text files the PRA searches on demand with its tools, an approach the authors found worked better than appending turns to its context. The PRA's checks run in its own writable scratch space; the base agent's workspace is read-only to it.

Rubrics are discovered offline from past harness-only runs by contrasting stronger and weaker attempts on the same task (top versus bottom quarter), combining trajectory evidence with workspace artifacts, and clustering into a MECE set. The resulting 19 seed rubrics are used unchanged across configurations and applied as-is to DSPredict-Hard. The PRA then adapts them per task through add, update, and retire tools in an isolated per-task store, where adding requires having run and recorded a check for the rubric in the same review.

Evaluation reports four task-level rates—mean leaderboard percentile, valid submission rate, above-median rate, and medal rate—over a fixed denominator of every task, counting invalid or missing submissions as zero percentile. Selection policies include first attempt, last attempt, mean over attempts, and best attempt (an oracle upper bound), with significance tested by a one-sided paired permutation test over tasks.

Why This Matters

Impact on research. The paper reframes the data science agent problem from "generate a better pipeline" to "supervise the process and stop wasting repeated attempts." It provides evidence that rubric-based supervision, mined automatically and adapted per task, outperforms a strong step-level PRM with code execution (DataPRM), and that process failures are properties of the agent rather than the domain—which is why a rubric set mined on tabular Kaggle tasks transfers to audio, vision, and molecular tasks.

Real-world applications:

• Automating Kaggle-style predictive modeling competitions and internal leaderboard challenges where repeated submissions are allowed.
• Building ML engineering assistants that work within a fixed compute budget and must avoid re-running the same baselines.
• Auditing data science pipelines for silent failures such as leaked validation sets, invalid cross-validation, or submissions that don't match the reported pipeline.
• Domains where prediction errors are quiet and costly—speech command classification (the paper's 12-command, 158,538-clip example where a wrong file path made every clip load as silent audio), earthquake time-to-failure prediction, and molecular property prediction.

Industry relevance. Teams deploying coding agents on data work need exactly the failure modes named here—budget waste, unsound validation, submission integrity—handled by an automated auditor that runs code rather than trusting the trajectory's own claims. The finding that the PRA can consume 75% of a run's cost on hard tasks also gives practitioners a concrete cost/benefit benchmark for supervision cadence.

Future Directions

1. Train the base agent on PRA feedback so it internalizes the process criteria and depends less on supervision at test time.
2. Credit assignment over rubrics. The PRA never retired a rubric and made only one update in 177 edits; measuring each rubric's effect on outcomes could make adaptation more selective than a store that only grows.
3. Smaller open-weight PRAs to replace the closed-source Claude Sonnet 4.6 supervisor, extending prior work from non-agentic critics to agentic process reward agents—also addressing the reproducibility limitation that exact results depend on a commercial model remaining available and unchanged.
4. Harness evolution (the paper cites Meta-Harness) and better selection among attempts, since all reported multi-attempt results use exactly five attempts and the paper does not study how gains change as attempts grow.

Target Audience

Researchers and engineers working on LLM agents, agent evaluation, and reward modeling—particularly those building coding or data-science agents that run for hundreds of turns. It is also relevant to practitioners who deploy agents on ML pipelines and need automated auditing for silent process errors, and to readers interested in automatic rubric or checklist construction as an alternative to scalar reward models. Readers without background in agents and process reward models will find the core ideas intuitive but the experimental setup dense.

Authors’ abstract

Data science agents can automate end-to-end machine learning workflows, yet they remain far from competitive on human leaderboards. Agents make process errors such as insufficient exploration, unsound validation, coding bugs, and poor use of a fixed budget (re-running the same exploratory data analysis and baseline models in every attempt). We propose AdaPRA to address these issues: a persistent coding agent harness that accumulates knowledge across attempts on the same task, with rubric-guided process supervision that identifies and corrects errors. Our process reward agent audits the agent's trajectory and workspace by executing checks, and adapts its rubrics to each task. We find that rubrics seeded from easier tasks generalize to harder tasks, with rubric adaptation further improving performance. With five attempts per task, AdaPRA improves an 80B Qwen model's leaderboard percentile by 24.5 points on DSPredict-Easy and 14.8 points on DSPredict-Hard while reducing process errors. We also show that supervision with rubrics mined from one base model's runs improves another (DeepSeek-V3.2), and that AdaPRA beats DataPRM, a competitive PRM for data science. We also share our code for reproducibility.