# Teaching an LLM to Fix Data — Without Touching the Model

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

🤗 [Live Environment](https://huggingface.co/spaces/Aswini-Kumar/data-centric-env) · 💻 [GitHub](https://github.com/CelestialWorthyOfHeavenAndEarth/data-centric-env) · 📓 [Training Notebook](https://colab.research.google.com/github/CelestialWorthyOfHeavenAndEarth/data-centric-env/blob/main/train_colab.ipynb)

---

## 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](https://datacentricai.org/) 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](https://github.com/meta-pytorch/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](https://huggingface.co/docs/trl/grpo_trainer) 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):
```python
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

![GRPO training reward over 150 episodes](plots/reward_curve.png)

*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.*

![Training dashboard: success rate per level, accuracy gain, curriculum progression](plots/training_dashboard.png)

### 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%** |

![Comparison bar chart: Random vs Heuristic vs Trained agent success rates](plots/baseline_comparison.png)

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.*