BLOG.md

Teaching an LLM to Fix Data — Without Touching the Model

OpenEnv Hackathon 2026 | Theme #3.1: World Modeling / Professional Tasks

🤗 Live Environment · 💻 GitHub · 📓 Training Notebook

---

The Problem

In ML, we spend 80% of our time cleaning data — yet almost no training infrastructure exists to teach LLMs to do it systematically.

Andrew Ng's Data-Centric AI movement shows that fixing the data outperforms tuning the model. Better labels, fewer missing values, and balanced class distributions consistently beat bigger architectures on real-world tasks. Yet the field has no reinforcement learning environment for training this skill.

We built one.

---

What We Built

Data-Centric AI Agent is an OpenEnv-compliant RL environment where a language model learns to orchestrate specialist data-cleaning sub-agents.

The setup is deliberately constrained — the ML classifier is frozen. The agent cannot touch it. Its only lever is fixing the data. It wins by pushing classifier accuracy past a target threshold within a step budget.

This forces genuine data engineering reasoning: what's broken, which specialist knows how to fix it, in what order.

---

The Environment

Multi-Agent Architecture

The LLM doesn't clean data directly. It talks to four specialist sub-agents and decides which recommendations to apply:

Specialist	Capability
CleanerAgent	Missing values, outliers, type errors — selects fill strategy via skewness analysis
AugmenterAgent	Synthesizes rows for underrepresented classes using Gaussian noise
BalancerAgent	Recommends oversample/undersample for class imbalance
ValidatorAgent	Detects rule violations and cross-column inconsistencies
AnalystAgent	Meta-specialist: holistic diagnosis + prioritised fix plan

A typical agent interaction:

query_analyst     → "Biggest issue: 23% missing in feature_2 (skewed). Secondary: class 1 at 18%"
query_cleaner     → "[1] Fill feature_2 with median (impact: +0.09, confidence: 0.91)"
apply 1           → Applied. 0 missing remaining. Row count: 200/200
validate          → RF Accuracy: 0.71 (+0.08). Target: 0.79. Not yet.
query_balancer    → "[1] Oversample class 1 to 80 samples (impact: +0.07)"
apply 1           → Applied.
submit            → Final accuracy: 0.81. TARGET HIT ✓  Reward: +0.74

Reward — Discriminating, Not Trivial

The original design gave a flat +0.50 bonus whenever the agent hit the target — regardless of how it got there. 5 steps or 25 steps got the same reward. This caused immediate saturation: the model hit 100% success rate in 10 training steps with no gradient to improve.

We redesigned the reward from scratch around a core principle: reward must discriminate between good strategies and mediocre ones.

The new AccuracyRubric at submit:

reward = 0.35 (base)
       + 0.30 × (budget_remaining / budget_total)  ← efficiency multiplier
       + min(0.15, (current - target) × 3.0)       ← stretch bonus for exceeding target

An agent that hits the target in 5 of 30 steps earns +0.79. One that barely scrapes it in 28 steps earns +0.37. Failing to hit target: up to −0.40. This creates a real learning gradient.

4-Level Curriculum

Level	Dataset	Issues	Budget	Target
Easy	200 rows	Missing + imbalance	25 steps	0.79 accuracy
Medium	500 rows	Missing + duplicates + imbalance + type errors	40 steps	0.74 accuracy
Hard	900 rows	6 issue types incl. outliers	60 steps	0.71 accuracy

Training starts at Easy (tutorial skipped — too trivial). Advancement criterion: ≥70% success rate over 20 episodes.

---

Training Pipeline

We train Qwen2.5-1.5B-Instruct (4-bit QLoRA via Unsloth, r=8) — chosen for fast iteration: ~3× faster per step than 3B, fits comfortably in 16GB VRAM, and still highly capable.

Phase 1: SFT Warmup (~15 min)

The model starts with no knowledge of our command format. Without warmup it outputs free text and gets zero reward forever. We run 1 epoch of supervised fine-tuning on 9,480 heuristic trajectory examples generated by sft_generator.py.

Phase 2: GRPO Training (~45–90 min)

TRL's GRPOTrainer drives live environment rollouts:

1. Model generates a command
2. Command sent to the live WebSocket environment server
3. Environment executes it, returns the rubric-graded reward
4. Gradients flow back through GRPO

Key hyperparameters (tuned for fast iteration):

per_device_train_batch_size = 2
gradient_accumulation_steps = 2
num_generations = 2          # rollouts per step
max_completion_length = 30   # commands are short

Experiment tracking: TensorBoard (logs/sft/ and logs/grpo/). No external API needed.

---

Results

Training Curves

After 150 GRPO episodes:

• Mean reward rose from −0.10 (random) to +0.65 (trained)
• Success rate on Easy: 30% → 95%
• Success rate on Medium: 0% → 80%
• Max curriculum level reached: Hard (level 3) at episode 110

The rolling mean (blue line) rises steadily through 3 curriculum levels. The sharp initial dip (episodes 5–15) is the model leaving the SFT distribution — a known GRPO warm-up effect.

Trained Agent vs Baselines

Agent	Tutorial	Easy	Medium	Hard	Overall
Random Agent	30%	20%	10%	5%	16%
Heuristic Baseline	100%	80%	60%	40%	70%
Trained Agent (GRPO)	100%	95%	80%	55%	83%

The trained agent outperforms the heuristic on all tasks. Crucially, the heuristic uses a fixed sequence (inspect → clean → balance → submit) that ignores actual data issues. The trained agent reads the data and adapts: on a dataset with no missing values but severe imbalance, it goes directly to query_balancer rather than wasting budget on query_cleaner.

Observed Behaviour Change

Before training, typical agent trajectory:

apply 1            ← blind apply (no query) → −0.08 penalty
validate           ← reward 0.0 (cached)
validate           ← redundant → −0.08 penalty
submit             ← target missed → −0.40 penalty
Total episode reward: −0.38

After GRPO training:

query_analyst      ← diagnosis first → correct workflow
query_cleaner
apply 1            ← after query → +0.05 process bonus
validate           ← after apply → +0.04 process bonus
query_balancer
apply 1
submit             ← target hit efficiently → +0.74 total
Total episode reward: +0.74

The model learned the correct sequencing entirely from reward signal — not from hardcoded rules.

---

Why It Matters

Data-Centric AI is the underexplored frontier of LLM capabilities. Every production ML system has a data quality problem. An LLM that can autonomously diagnose and fix those problems — adapting its strategy to the actual data, not a fixed playbook — is a genuinely useful tool that doesn't exist yet.

This environment provides the training ground for that capability.

Built for the OpenEnv Hackathon 2026. Theme #3.1: World Modeling / Professional Tasks.