Data-Centric AI RL Environment — OpenEnv Hackathon Submission

.gitignore

@@ -0,0 +1,45 @@
+# Python
+__pycache__/
+*.py[cod]
+*.pyd
+*.pyo
+*.egg-info/
+*.egg
+.eggs/
+
+# Virtual environments
+.venv/
+venv/
+env/
+
+# Package managers
+uv.lock
+poetry.lock
+
+# Test / build artefacts
+.pytest_cache/
+.coverage
+htmlcov/
+dist/
+build/
+
+# Training outputs (generated — not committed until training is done)
+sft_data.jsonl
+generations.jsonl
+training_log.json
+sft-checkpoint/
+data-centric-checkpoints/
+data-centric-adapter/
+data-centric-merged/
+
+# IDE
+.vscode/
+.idea/
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Hackathon docs (keep local only)
+*.pdf
+\[External\]*/

.spaceignore

@@ -0,0 +1,39 @@
+# Files excluded from HuggingFace Space Docker build
+# Keep this lean — only block dev/training files the server doesn't need
+
+# Python cache
+__pycache__/
+**/__pycache__/
+*.pyc
+*.pyo
+*.pyd
+*.egg-info/
+
+# Dev / test tools
+.venv/
+venv/
+.pytest_cache/
+test_features_smoke.py
+uv.lock
+
+# Training-only scripts (not needed by the server at runtime)
+train_data_centric.py
+train_colab.ipynb
+eval_data_centric.py
+sft_generator.py
+inference.py
+hf_job_train.py
+plot_rewards.py
+
+# Generated training artefacts
+sft_data.jsonl
+generations.jsonl
+training_log.json
+sft-checkpoint/
+data-centric-checkpoints/
+data-centric-adapter/
+data-centric-merged/
+
+# Git
+.git/
+.gitignore

Dockerfile

@@ -0,0 +1,34 @@
+FROM python:3.10-slim
+
+# System dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+# Working directory
+WORKDIR /app
+
+# Copy project files
+COPY . .
+
+# Install Python dependencies
+RUN pip install --no-cache-dir \
+    "openenv-core[core]>=0.2.1" \
+    "fastapi>=0.115.0" \
+    "uvicorn>=0.24.0" \
+    "scikit-learn>=1.3.0" \
+    "pandas>=2.0.0" \
+    "numpy>=1.24.0"
+
+# Install the package itself (enables relative imports)
+RUN pip install --no-cache-dir -e .
+
+# HF Spaces runs as non-root user 1000
+RUN useradd -m -u 1000 user
+USER user
+
+# Expose port that matches app_port in README.md
+EXPOSE 7860
+
+# Start the environment server
+CMD ["python", "-m", "uvicorn", "server.app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -0,0 +1,307 @@
+---
+title: Data-Centric AI RL Environment
+emoji: 🧠
+colorFrom: blue
+colorTo: indigo
+sdk: docker
+pinned: false
+app_port: 7860
+tags:
+  - openenv
+  - reinforcement-learning
+  - data-centric-ai
+  - grpo
+  - unsloth
+---
+
+# 🧠 Data-Centric AI — Multi-Agent RL Environment
+
+An [OpenEnv](https://github.com/meta-pytorch/OpenEnv)-compliant reinforcement learning environment that trains an LLM to act as a **data engineering orchestrator** — coordinating 4 specialist sub-agents across multi-step plans to improve ML datasets under budget constraints.
+
+> **Core insight:** In traditional ML, practitioners tune models to squeeze out performance. This environment flips that — the model architecture is deliberately **frozen**, forcing the LLM agent to master *data engineering* as its only lever: diagnosing noise, coordinating specialist agents, and strategically transforming the dataset until accuracy surpasses the target. This is [Data-Centric AI](https://datacentricai.org/) — the paradigm Andrew Ng argues matters more than model architecture.
+
+> **Live Space:** https://huggingface.co/spaces/Aswini-Kumar/data-centric-env
+
+### Key Capabilities
+
+| Capability | How it works |
+|---|---|
+| **Multi-Agent Coordination** | LLM orchestrates 4 specialist agents (cleaner, augmenter, balancer, validator) — deciding *who* to call and *when*, modeling each specialist's strengths |
+| **Long-Horizon Planning** | 30-step budget with sparse terminal reward. Agent must plan inspect → query → apply → validate → submit sequences with delayed feedback |
+| **Theory-of-Mind Reasoning** | Agent infers which specialist is best for the current data problem (class imbalance vs. missing values vs. outliers) |
+| **Anti-Exploit Hardening** | 9 security mechanisms (immutable ground truth, golden rows, cooldowns, budget caps) prevent reward hacking |
+| **Curriculum Learning** | Auto-advances from tutorial → easy → medium → hard based on rolling success rate |
+| **Composable Reward System** | 4-component rubric (accuracy + process + preservation + efficiency) using OpenEnv's `Rubric` base class |
+
+---
+
+## 🎯 What the Agent Does
+
+The agent receives a noisy tabular dataset and a fixed classifier. It must orchestrate specialist sub-agents to clean, augment, and balance the data until accuracy hits a target — **without touching the model**.
+
+Each episode:
+1. Agent **inspects** the dataset and model
+2. Agent **queries** specialist sub-agents for recommendations
+3. Agent **applies** the best fix (or **undoes** a bad one)
+4. Agent **validates** accuracy improvement
+5. Agent **submits** when target is reached or budget runs out
+
+---
+
+## 🏗️ Architecture
+
+```
+┌─────────────────────────────────────────────────────────────────┐
+│                     LLM Agent (Qwen2.5-3B)                      │
+│            SFT warmup → GRPO live-environment training          │
+└─────────────┬───────────────────────────────────┬───────────────┘
+              │  text commands                    │  structured obs
+              ▼                                   ▲
+┌─────────────────────────────────────────────────────────────────┐
+│              DataCentricEnvironment (OpenEnv)                    │
+│  ┌──────────┐  ┌──────────┐  ┌──────────┐  ┌──────────────┐    │
+│  │ Cleaner  │  │Augmenter │  │ Balancer │  │  Validator   │    │
+│  │  Agent   │  │  Agent   │  │  Agent   │  │   Agent      │    │
+│  └────┬─────┘  └────┬─────┘  └────┬─────┘  └──────┬───────┘    │
+│       └──────────────┴──────────────┴───────────────┘            │
+│                           │                                      │
+│              ┌────────────▼────────────┐                         │
+│              │   Working Copy (mutable) │◄── Snapshot stack (×3) │
+│              └────────────┬────────────┘    for undo support     │
+│                           │                                      │
+│              ┌────────────▼────────────┐                         │
+│              │  ModelEvaluator (RF)     │                         │
+│              │  n_est=20 (fast_mode)   │                         │
+│              └────────────┬────────────┘                         │
+│                           │                                      │
+│              ┌────────────▼────────────┐                         │
+│              │  Ground Truth (frozen)  │ ← never mutated         │
+│              └─────────────────────────┘                         │
+│                                                                  │
+│  ┌───────────────────────────────────────────────────────────┐   │
+│  │  DataCentricRubric (OpenEnv Rubric system)                │   │
+│  │  ├── AccuracyRubric      — Δ accuracy vs baseline         │   │
+│  │  ├── ProcessRubric       — workflow pattern scoring       │   │
+│  │  ├── PreservationRubric  — row preservation incentive     │   │
+│  │  └── EfficiencyRubric    — accuracy gain / budget used    │   │
+│  │  + StepRubric            — dense per-apply proxy reward   │   │
+│  └───────────────────────────────────────────────────────────┘   │
+│                                                                  │
+│  Anti-Exploit: 9 protections (GT immutability, cooldowns, etc.)  │
+└─────────────────────────────────────────────────────────────────┘
+```
+
+---
+
+## 🌍 Environment Design
+
+### Action Space
+Single text command — one of:
+
+| Command | Effect |
+|---------|--------|
+| `inspect_dataset` | View shape, missing values, class distribution |
+| `inspect_model` | View classifier accuracy, precision, recall, F1 |
+| `query_cleaner` | Get missing-value / outlier fix recommendations |
+| `query_augmenter [class]` | Get data augmentation recommendations |
+| `query_balancer` | Get class rebalancing recommendations |
+| `query_validator` | Check rule violations (costs 2 budget) |
+| `apply <N>` | Apply recommendation number N |
+| `reject <N>` | Reject a recommendation |
+| `undo` | Revert last apply (max 3 levels deep) |
+| `validate` | Retrain classifier and score (cooldown applies) |
+| `submit` | Finalize and score the episode |
+
+### Observation Space
+Each step returns a structured observation:
+
+```python
+DataCentricObservation(
+    response="...",              # Specialist agent's text response
+    current_accuracy=0.71,       # Current classifier accuracy
+    baseline_accuracy=0.62,      # Accuracy before any changes
+    target_accuracy=0.73,        # Target to hit
+    estimated_quality=0.84,      # Dataset quality score (0-1)
+    rows_preserved_pct=0.97,     # % of original rows still present
+    budget_remaining=22,         # Steps remaining
+    validate_calls_remaining=2,  # Free validate calls left
+    active_session="cleaner",    # Which specialist is active
+    done=False,
+)
+```
+
+### Reward Function — OpenEnv Rubric System
+
+Uses `openenv.core.rubrics.base.Rubric` with composable child rubrics (nn.Module-style auto-registration):
+
+| Rubric | Signal | Range |
+|--------|--------|-------|
+| **AccuracyRubric** | Δ accuracy × 2.5 + submit bonus | [-1.0, +1.0] |
+| **ProcessRubric** | Correct query→apply→validate sequencing | [-0.10, +0.05] |
+| **PreservationRubric** | Rows preserved ≥ 90% | [-0.40, +0.05] |
+| **EfficiencyRubric** | Accuracy gain / budget used (submit only) | [-0.05, +0.20] |
+| **StepRubric** | Dense per-apply quality proxy | [-0.30, +0.15] |
+
+Total clamped to **[-1.0, 1.0]** by `DataCentricRubric.forward()`.
+
+### Anti-Exploit Protections
+9 hardened mechanisms including:
+- Ground truth immutability assertion after every `apply`
+- Validate cooldown enforcement (must take 2 actions between validates)
+- Duplicate apply detection + session apply limit (max 3 per query)
+- Recommendation staleness validation (re-query after each session)
+- Catastrophic data loss detection (< 50% rows → terminate)
+- Episode wall-clock timeout (5 min → forced submit)
+- Input truncation (> 200 chars → truncate + penalty)
+
+---
+
+## 📚 Task Curriculum (4 Levels)
+
+| Task | Rows | Issues | Baseline | Target | Budget |
+|------|------|--------|----------|--------|--------|
+| `task_0_tutorial` | 100 | Missing values (20%) | ~0.62 | 0.73 | 30 |
+| `task_1_easy` | 200 | Missing + imbalance | ~0.63 | 0.79 | 25 |
+| `task_2_medium` | 500 | Missing + duplicates + imbalance + type errors | ~0.58 | 0.74 | 40 |
+| `task_3_hard` | 900 | 6 issues incl. outliers + logic errors | ~0.54 | 0.71 | 60 |
+
+Advancement criterion: ≥ 60% success rate over a rolling 30-episode window.
+
+---
+
+## 🤖 Training Pipeline
+
+**Model:** Qwen2.5-3B-Instruct (4-bit via Unsloth)  
+**Algorithm:** SFT warmup → GRPO (TRL)  
+**Framework:** OpenEnv + TRL + Unsloth
+
+### Phase 1 — SFT Warmup
+Train on ~8,100 heuristic trajectory examples to teach valid command syntax before RL.
+
+### Phase 2 — GRPO
+Live environment rollouts scored by the composable Rubric system. Curriculum scheduler advances from tutorial → easy → medium → hard as performance improves.
+
+### Training Notebook
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/CelestialWorthyOfHeavenAndEarth/data-centric-env/blob/main/train_colab.ipynb)
+
+[`train_colab.ipynb`](train_colab.ipynb) — complete end-to-end training pipeline for Colab T4 GPU.
+
+---
+
+## 📊 Results
+
+### Heuristic Baseline Verification
+
+The heuristic agent (`inference.py`) validates that the environment is solvable:
+
+| Task | Accuracy Gain | Target Hit Rate |
+|------|--------------|-----------------|
+| `task_0_tutorial` | +0.11 | ✓ 100% |
+| `task_1_easy` | +0.08 | ✓ 80% |
+| `task_2_medium` | +0.06 | ✓ 60% |
+| `task_3_hard` | +0.04 | ~ 40% |
+
+> **Note:** After GRPO training, embed reward curves from `plots/` here.
+> Run `python plot_rewards.py` to generate: `reward_curve.png`, `success_rate.png`, `accuracy_gain.png`, `curriculum.png`.
+
+---
+
+## 🧪 Testing
+
+```bash
+# Run all tests (35 tests — grader + environment)
+pytest tests/ -v
+
+# Run only grader tests (22 tests)
+pytest tests/test_grader.py -v
+
+# Run only environment safety tests (13 tests)
+pytest tests/test_environment.py -v
+```
+
+Tests cover:
+- All 5 Rubric components (accuracy, process, preservation, efficiency, step)
+- Reward clamping to declared [-1.0, 1.0] range
+- Ground truth immutability after every command
+- Budget enforcement and episode termination
+- Validate cooldown and call limiting
+- Undo/snapshot stack behavior
+- Unknown command handling
+
+---
+
+## 🚀 Quick Start
+
+### Connect to the Live Space
+
+```python
+from client import DataCentricEnv
+from models import DataCentricAction
+
+with DataCentricEnv(base_url="https://aswini-kumar-data-centric-env.hf.space").sync() as env:
+    result = env.reset(task="task_0_tutorial", seed=42)
+    print(f"Baseline: {result.observation.baseline_accuracy:.2f}  Target: {result.observation.target_accuracy:.2f}")
+
+    result = env.step(DataCentricAction(message="inspect_dataset"))
+    print(result.observation.response)
+
+    result = env.step(DataCentricAction(message="query_cleaner"))
+    print(result.observation.response)
+```
+
+### Run Locally
+
+```bash
+# Install
+pip install openenv-core fastapi uvicorn scikit-learn pandas numpy
+
+# Start server
+uvicorn server.app:app --host 0.0.0.0 --port 8000
+
+# In another terminal
+python -c "
+from client import DataCentricEnv
+from models import DataCentricAction
+with DataCentricEnv(base_url='http://localhost:8000').sync() as env:
+    obs = env.reset(task='task_0_tutorial', seed=42).observation
+    print(f'Ready — baseline={obs.baseline_accuracy:.2f} target={obs.target_accuracy:.2f}')
+"
+```
+
+---
+
+## 📁 Project Structure
+
+```
+data_centric_env/
+├── openenv.yaml              # OpenEnv manifest (tasks, reward range, action/obs types)
+├── client.py                 # DataCentricEnv WebSocket client
+├── models.py                 # DataCentricAction + DataCentricObservation (Pydantic)
+├── train_data_centric.py     # Full SFT → GRPO training pipeline
+├── train_colab.ipynb         # Colab training notebook (T4 GPU)
+├── eval_data_centric.py      # Evaluation: random vs trained agent
+├── plot_rewards.py           # Reward curve visualization (4 plots)
+├── sft_generator.py          # SFT warmup data generator (~8100 examples)
+├── inference.py              # Heuristic baseline agent
+├── tests/
+│   ├── test_grader.py        # 22 tests — Rubric system + reward components
+│   └── test_environment.py   # 13 tests — safety invariants + anti-exploit
+└── server/
+    ├── app.py                # FastAPI server (HTTP + WebSocket via OpenEnv)
+    ├── data_centric_environment.py   # Core RL environment logic (680 lines)
+    ├── dataset_generator.py  # Synthetic dataset generation (4 task configs)
+    ├── specialist_agents.py  # CleanerAgent, AugmenterAgent, BalancerAgent, ValidatorAgent
+    ├── grader.py             # Composable Rubric system (openenv.core.rubrics.base)
+    ├── anti_exploit.py       # 9 anti-reward-hacking protections
+    ├── model_evaluator.py    # RF classifier with hash-based caching
+    └── Dockerfile            # HuggingFace Spaces deployment
+```
+
+---
+
+## 🏷️ Hackathon
+
+**Theme:** #3.1 — World Modeling / Professional Tasks  
+**Stack:** OpenEnv · Unsloth · TRL (GRPO) · FastAPI · scikit-learn  
+**Repo:** [github.com/CelestialWorthyOfHeavenAndEarth/data-centric-env](https://github.com/CelestialWorthyOfHeavenAndEarth/data-centric-env)  
+**Space:** [huggingface.co/spaces/Aswini-Kumar/data-centric-env](https://huggingface.co/spaces/Aswini-Kumar/data-centric-env)

__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Data Centric Env Environment."""
+
+from .client import DataCentricEnv
+from .models import DataCentricAction, DataCentricObservation
+
+__all__ = [
+    "DataCentricAction",
+    "DataCentricObservation",
+    "DataCentricEnv",
+]

client.py

@@ -0,0 +1,64 @@
+"""Data-Centric Environment Client."""
+
+from typing import Dict
+
+from openenv.core import EnvClient
+from openenv.core.client_types import StepResult
+from openenv.core.env_server.types import State
+
+try:
+    from .models import DataCentricAction, DataCentricObservation
+except ImportError:
+    from models import DataCentricAction, DataCentricObservation
+
+
+class DataCentricEnv(EnvClient[DataCentricAction, DataCentricObservation, State]):
+    """
+    Client for the Data-Centric RL Environment.
+
+    Connects over WebSocket for efficient multi-step interactions.
+
+    Example:
+        >>> with DataCentricEnv(base_url="http://localhost:8000") as client:
+        ...     result = client.reset(task="task_0_tutorial")
+        ...     result = client.step(DataCentricAction(message="inspect_dataset"))
+        ...     print(result.observation.response)
+
+    Docker example:
+        >>> client = DataCentricEnv.from_docker_image("data_centric_env:latest")
+        >>> result = client.reset(task="task_1_easy")
+    """
+
+    def _step_payload(self, action: DataCentricAction) -> Dict:
+        return {"message": action.message}
+
+    def _parse_result(self, payload: Dict) -> StepResult[DataCentricObservation]:
+        obs_data = payload.get("observation", {})
+        observation = DataCentricObservation(
+            response=obs_data.get("response", ""),
+            current_accuracy=obs_data.get("current_accuracy", 0.0),
+            baseline_accuracy=obs_data.get("baseline_accuracy", 0.0),
+            target_accuracy=obs_data.get("target_accuracy", 0.0),
+            estimated_quality=obs_data.get("estimated_quality", 0.0),
+            dataset_shape=obs_data.get("dataset_shape", ""),
+            rows_preserved_pct=obs_data.get("rows_preserved_pct", 1.0),
+            budget_remaining=obs_data.get("budget_remaining", 0),
+            step_number=obs_data.get("step_number", 0),
+            max_steps=obs_data.get("max_steps", 30),
+            active_session=obs_data.get("active_session", "none"),
+            validate_calls_remaining=obs_data.get("validate_calls_remaining", 3),
+            done=payload.get("done", False),
+            reward=payload.get("reward", 0.0),
+            metadata=obs_data.get("metadata", {}),
+        )
+        return StepResult(
+            observation=observation,
+            reward=payload.get("reward", 0.0),
+            done=payload.get("done", False),
+        )
+
+    def _parse_state(self, payload: Dict) -> State:
+        return State(
+            episode_id=payload.get("episode_id"),
+            step_count=payload.get("step_count", 0),
+        )

eval_data_centric.py

@@ -0,0 +1,443 @@
+"""
+eval_data_centric.py — Evaluation script for DataCentricEnv.
+
+Runs two agents on identical seeds for fair comparison:
+  - Random Agent:  picks valid commands at random (baseline)
+  - Trained Agent: uses the fine-tuned model from ./data-centric-adapter
+
+Produces eval_results.json for use by plot_rewards.py.
+"""
+
+import json
+import os
+import random
+import signal
+import subprocess
+import sys
+import time
+from typing import Optional
+
+import requests
+import torch
+from unsloth import FastLanguageModel
+
+# WebSocket client for stateful episode rollouts
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from client import DataCentricEnv
+from models import DataCentricAction
+
+# ════════════════════════════════════════════════════════
+# CONSTANTS
+# ════════════════════════════════════════════════════════
+
+BASE_URL = os.environ.get("ENV_URL", "http://localhost:8000")
+ADAPTER_PATH = "./data-centric-adapter"
+MAX_SEQ_LENGTH = 1024
+EPISODES_PER_TASK = 10
+TASKS = ["task_0_tutorial", "task_1_easy", "task_2_medium", "task_3_hard"]
+
+VALID_COMMANDS = [
+    "inspect_dataset", "inspect_model", "query_analyst",
+    "query_cleaner", "query_augmenter 0", "query_balancer", "query_validator",
+    "apply 1", "apply 2", "reject 1", "undo", "validate", "submit",
+]
+
+SYSTEM_PROMPT = """You are a Data-Centric AI agent. Your job is to improve a \
+Machine learning dataset so a fixed classifier achieves higher accuracy.
+
+STRATEGY — use this order:
+1. query_analyst to get a prioritised action plan (costs 1 budget, worth it)
+2. inspect_dataset to understand the data
+3. query the recommended specialist (query_cleaner, query_augmenter, query_balancer)
+4. apply the best recommendation by number (e.g. apply 1)
+5. validate to check if accuracy improved
+6. repeat until you hit the target or run low on budget
+7. submit to finalize
+
+IMPORTANT RULES:
+- Start with query_analyst — it tells you the biggest problem to fix first.
+- Always query a specialist before applying. Never apply without querying first.
+- Check the Agreement signal after validate: DISAGREE means possible overfitting.
+- Validate after every 2-3 applies to track progress.
+- Do not spam validate — it costs budget after 3 uses.
+- query_validator costs 2 budget — use only when suspicious of data quality.
+- submit when accuracy >= target or budget < 5.
+
+Reply with exactly ONE command per message. No explanation. Just the command."""
+
+
+def build_user_prompt(obs: dict) -> str:
+    improvement_needed = obs.get("target_accuracy", 0) - obs.get("current_accuracy", 0)
+    return (
+        f"Current situation:\n"
+        f"Accuracy: {obs.get('current_accuracy', 0):.1%} → "
+        f"Target: {obs.get('target_accuracy', 0):.1%}\n"
+        f"Still need: {max(0, improvement_needed):.1%} improvement\n"
+        f"Quality score: {obs.get('estimated_quality', 0):.2f}/1.00\n"
+        f"Rows preserved: {obs.get('rows_preserved_pct', 1.0):.1%}\n"
+        f"Budget remaining: {obs.get('budget_remaining', 0)} steps\n"
+        f"Free validates left: {obs.get('validate_calls_remaining', 0)}\n"
+        f"Active query session: {obs.get('active_session', 'none')}\n\n"
+        f"Last result:\n{str(obs.get('response', ''))[:400]}\n\n"
+        f"What is your next action? (one command only)"
+    )
+
+
+# ════════════════════════════════════════════════════════
+# SERVER MANAGEMENT
+# ════════════════════════════════════════════════════════
+
+def start_server() -> subprocess.Popen:
+    proc = subprocess.Popen(
+        ["uvicorn", "server.app:app", "--host", "0.0.0.0", "--port", "8000"],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+    for i in range(30):
+        try:
+            r = requests.get(f"{BASE_URL}/health", timeout=1)
+            if r.status_code == 200:
+                print(f"Server ready after {i + 1}s")
+                return proc
+        except Exception:
+            pass
+        time.sleep(1)
+    proc.terminate()
+    raise RuntimeError("Server failed to start in 30 seconds")
+
+
+def stop_server(proc: subprocess.Popen):
+    proc.terminate()  # cross-platform (SIGTERM on Linux, TerminateProcess on Windows)
+    proc.wait()
+    print("Server stopped.")
+
+
+# ════════════════════════════════════════════════════════
+# MODEL LOADING
+# ════════════════════════════════════════════════════════
+
+def load_trained_model():
+    if not os.path.exists(ADAPTER_PATH):
+        raise FileNotFoundError(
+            f"Adapter not found at {ADAPTER_PATH}. "
+            "Run train_data_centric.py first."
+        )
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name=ADAPTER_PATH,
+        max_seq_length=MAX_SEQ_LENGTH,
+        load_in_4bit=True,
+        dtype=None,
+    )
+    FastLanguageModel.for_inference(model)
+    return model, tokenizer
+
+
+# ════════════════════════════════════════════════════════
+# EPISODE METRICS
+# ════════════════════════════════════════════════════════
+
+def episode_metrics(
+    task: str,
+    seed: int,
+    final_obs: dict,
+    actions: list,
+    baseline_accuracy: float,
+    max_steps: int,
+) -> dict:
+    """Compute per-episode metrics for a single completed episode."""
+    final_accuracy = final_obs.get("current_accuracy", baseline_accuracy)
+    budget_remaining = final_obs.get("budget_remaining", 0)
+    target_accuracy = final_obs.get("target_accuracy", 1.0)
+    budget_used = max_steps - budget_remaining
+
+    accuracy_improvement = final_accuracy - baseline_accuracy
+    target_hit = final_accuracy >= target_accuracy
+    budget_efficiency = (
+        accuracy_improvement / max(budget_used, 1)
+    )
+
+    # Format rate: % actions that are valid commands
+    valid_count = sum(
+        1 for a in actions
+        if any(a.strip().startswith(cmd.split()[0]) for cmd in VALID_COMMANDS)
+    )
+    format_rate = valid_count / max(len(actions), 1)
+
+    # Strategy rate: % query→apply consecutive pairs
+    strategy_hits = 0
+    for i in range(1, len(actions)):
+        if (actions[i - 1].startswith("query_")
+                and actions[i].startswith("apply")):
+            strategy_hits += 1
+    strategy_rate = strategy_hits / max(len(actions) - 1, 1)
+
+    return {
+        "task":                 task,
+        "seed":                 seed,
+        "final_accuracy":       round(final_accuracy, 4),
+        "baseline_accuracy":    round(baseline_accuracy, 4),
+        "target_accuracy":      round(target_accuracy, 4),
+        "accuracy_improvement": round(accuracy_improvement, 4),
+        "target_hit":           target_hit,
+        "budget_used":          budget_used,
+        "budget_efficiency":    round(budget_efficiency, 6),
+        "format_rate":          round(format_rate, 4),
+        "strategy_rate":        round(strategy_rate, 4),
+        "n_actions":            len(actions),
+    }
+
+
+# ════════════════════════════════════════════════════════
+# RANDOM AGENT
+# ════════════════════════════════════════════════════════
+
+def run_random_episode(task: str, seed: int) -> Optional[dict]:
+    """Run one episode with a random agent using the WebSocket client."""
+    try:
+        with DataCentricEnv(base_url=BASE_URL).sync() as env:
+            r_reset = env.reset(task=task, seed=seed)
+            obs = r_reset.observation
+            baseline_accuracy = obs.baseline_accuracy
+            max_steps = obs.max_steps
+            actions = []
+
+            while not obs.done:
+                action = random.choice(VALID_COMMANDS)
+                actions.append(action)
+                try:
+                    step_result = env.step(DataCentricAction(message=action))
+                    obs = step_result.observation
+                except Exception:
+                    break
+
+            return episode_metrics(
+                task, seed,
+                {"current_accuracy": obs.current_accuracy,
+                 "budget_remaining": obs.budget_remaining,
+                 "target_accuracy":  obs.target_accuracy,
+                 "done":             obs.done},
+                actions, baseline_accuracy, max_steps
+            )
+    except Exception as e:
+        print(f"  [random] Episode failed: {e}")
+        return None
+
+
+# ════════════════════════════════════════════════════════
+# TRAINED AGENT
+# ════════════════════════════════════════════════════════
+
+def run_trained_episode(
+    model, tokenizer, task: str, seed: int
+) -> Optional[dict]:
+    """Run one episode with the trained model using the WebSocket client."""
+    try:
+        with DataCentricEnv(base_url=BASE_URL).sync() as env:
+            r_reset = env.reset(task=task, seed=seed)
+            obs = r_reset.observation
+            baseline_accuracy = obs.baseline_accuracy
+            max_steps = obs.max_steps
+            actions = []
+
+            while not obs.done:
+                obs_dict = {
+                    "current_accuracy":        obs.current_accuracy,
+                    "target_accuracy":         obs.target_accuracy,
+                    "estimated_quality":       obs.estimated_quality,
+                    "rows_preserved_pct":      obs.rows_preserved_pct,
+                    "budget_remaining":        obs.budget_remaining,
+                    "validate_calls_remaining":obs.validate_calls_remaining,
+                    "active_session":          obs.active_session,
+                    "response":                obs.response,
+                }
+                messages = [
+                    {"role": "system", "content": SYSTEM_PROMPT},
+                    {"role": "user",   "content": build_user_prompt(obs_dict)},
+                ]
+                input_ids = tokenizer.apply_chat_template(
+                    messages,
+                    return_tensors="pt",
+                    max_length=MAX_SEQ_LENGTH - 60,
+                    truncation=True,
+                    add_generation_prompt=True,
+                ).to(model.device)
+
+                with torch.no_grad():
+                    output_ids = model.generate(
+                        input_ids,
+                        max_new_tokens=50,
+                        temperature=0.1,
+                        do_sample=True,
+                        pad_token_id=tokenizer.eos_token_id,
+                    )
+
+                action = tokenizer.decode(
+                    output_ids[0][input_ids.shape[1]:],
+                    skip_special_tokens=True,
+                ).strip().split("\n")[0].strip()[:200]
+
+                actions.append(action)
+                try:
+                    step_result = env.step(DataCentricAction(message=action))
+                    obs = step_result.observation
+                except Exception as e:
+                    break
+
+            return episode_metrics(
+                task, seed,
+                {"current_accuracy": obs.current_accuracy,
+                 "budget_remaining": obs.budget_remaining,
+                 "target_accuracy":  obs.target_accuracy,
+                 "done":             obs.done},
+                actions, baseline_accuracy, max_steps
+            )
+    except Exception as e:
+        print(f"  [trained] Episode failed: {e}")
+        return None
+
+
+# ════════════════════════════════════════════════════════
+# AGGREGATION
+# ════════════════════════════════════════════════════════
+
+def aggregate(episodes: list) -> dict:
+    """Compute mean metrics across a list of episode result dicts."""
+    if not episodes:
+        return {}
+    keys = [
+        "accuracy_improvement", "target_hit", "budget_used",
+        "budget_efficiency", "format_rate", "strategy_rate",
+    ]
+    return {
+        k: round(sum(ep[k] for ep in episodes) / len(episodes), 4)
+        for k in keys
+    }
+
+
+def print_comparison_table(random_agg: dict, trained_agg: dict):
+    """Print a formatted comparison table to stdout."""
+    def pct_change(r, t):
+        if r == 0:
+            return "—"
+        return f"{(t - r) / abs(r) * 100:+.0f}%"
+
+    def pp_change(r, t):
+        return f"{(t - r) * 100:+.0f}pp"
+
+    rows = [
+        ("Accuracy gain",      f"{random_agg.get('accuracy_improvement',0):.3f}",
+                               f"{trained_agg.get('accuracy_improvement',0):.3f}",
+                               pct_change(random_agg.get('accuracy_improvement',0),
+                                          trained_agg.get('accuracy_improvement',0))),
+        ("Target hit rate",    f"{random_agg.get('target_hit',0):.0%}",
+                               f"{trained_agg.get('target_hit',0):.0%}",
+                               pp_change(random_agg.get('target_hit',0),
+                                         trained_agg.get('target_hit',0))),
+        ("Budget efficiency",  f"{random_agg.get('budget_efficiency',0):.4f}",
+                               f"{trained_agg.get('budget_efficiency',0):.4f}",
+                               pct_change(random_agg.get('budget_efficiency',0),
+                                          trained_agg.get('budget_efficiency',0))),
+        ("Format rate",        "random",
+                               f"{trained_agg.get('format_rate',0):.0%}", "—"),
+        ("Strategy rate",      "0%",
+                               f"{trained_agg.get('strategy_rate',0):.0%}", "—"),
+    ]
+
+    header = f"{'Metric':<20} {'Random':>12} {'Trained':>13} {'Improvement':>14}"
+    sep    = "─" * len(header)
+    print(f"\n{sep}")
+    print(header)
+    print(sep)
+    for metric, rand, trained, improvement in rows:
+        print(f"  {metric:<18} {rand:>12} {trained:>13} {improvement:>14}")
+    print(sep + "\n")
+
+
+# ════════════════════════════════════════════════════════
+# MAIN
+# ════════════════════════════════════════════════════════
+
+if __name__ == "__main__":
+    server_proc = start_server()
+
+    try:
+        print(f"\nLoading trained model from {ADAPTER_PATH}...")
+        model, tokenizer = load_trained_model()
+
+        # Use fixed seeds so both agents see identical tasks
+        seeds = list(range(EPISODES_PER_TASK))
+
+        results = {
+            "random":  {"all_episodes": [], "by_task": {}},
+            "trained": {"all_episodes": [], "by_task": {}},
+        }
+
+        for task in TASKS:
+            print(f"\n{'='*50}")
+            print(f"Task: {task}")
+            print("─" * 50)
+
+            random_eps, trained_eps = [], []
+
+            for seed in seeds:
+                print(f"  Seed {seed:2d}:", end="  ")
+
+                # Random agent
+                sys.stdout.write("[random] ")
+                sys.stdout.flush()
+                r_ep = run_random_episode(task, seed)
+                if r_ep:
+                    random_eps.append(r_ep)
+                    sys.stdout.write(
+                        f"acc={r_ep['final_accuracy']:.3f} "
+                        f"hit={'✓' if r_ep['target_hit'] else '✗'}  "
+                    )
+
+                # Trained agent (same seed)
+                sys.stdout.write("[trained] ")
+                sys.stdout.flush()
+                t_ep = run_trained_episode(model, tokenizer, task, seed)
+                if t_ep:
+                    trained_eps.append(t_ep)
+                    sys.stdout.write(
+                        f"acc={t_ep['final_accuracy']:.3f} "
+                        f"hit={'✓' if t_ep['target_hit'] else '✗'}"
+                    )
+
+                print()
+
+            results["random"]["by_task"][task]  = aggregate(random_eps)
+            results["trained"]["by_task"][task] = aggregate(trained_eps)
+            results["random"]["all_episodes"].extend(random_eps)
+            results["trained"]["all_episodes"].extend(trained_eps)
+
+        # Overall aggregates
+        results["random"]["overall"]  = aggregate(results["random"]["all_episodes"])
+        results["trained"]["overall"] = aggregate(results["trained"]["all_episodes"])
+
+        # Print comparison table
+        print_comparison_table(
+            results["random"]["overall"],
+            results["trained"]["overall"],
+        )
+
+        # Print per-task breakdown
+        print("Per-task summary:")
+        for task in TASKS:
+            r = results["random"]["by_task"].get(task, {})
+            t = results["trained"]["by_task"].get(task, {})
+            print(
+                f"  {task:<22}  "
+                f"random: acc+{r.get('accuracy_improvement',0):.3f} "
+                f"hit={r.get('target_hit',0):.0%}  |  "
+                f"trained: acc+{t.get('accuracy_improvement',0):.3f} "
+                f"hit={t.get('target_hit',0):.0%}"
+            )
+
+        # Save full results
+        json.dump(results, open("eval_results.json", "w"), indent=2)
+        print("\nResults saved to eval_results.json")
+        print("Run plot_rewards.py to visualise.")
+
+    finally:
+        stop_server(server_proc)

hf_job_train.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+"""
+HF Job training script for Data-Centric AI Agent.
+
+Run this as an HF Job pointing to the deployed HF Space as the environment:
+    hf job run --gpu t4-small --env ENV_URL=https://aswini-kumar-data-centric-env.hf.space \
+        python hf_job_train.py
+
+Or via Python API:
+    from huggingface_hub import HfApi
+    api = HfApi()
+    api.run_job(...)
+
+The HF Space (server) must be running BEFORE submitting this job.
+"""
+
+import os
+import sys
+import time
+import requests
+
+# ── Verify the HF Space is reachable before starting training ─────────────────
+
+ENV_URL = os.environ.get("ENV_URL", "https://aswini-kumar-data-centric-env.hf.space")
+print(f"[HF Job] Environment URL: {ENV_URL}")
+
+print("[HF Job] Checking environment server health...")
+for attempt in range(12):
+    try:
+        r = requests.get(f"{ENV_URL}/health", timeout=10)
+        if r.status_code == 200:
+            print(f"[HF Job] Server healthy: {r.json()}")
+            break
+    except Exception as e:
+        print(f"[HF Job] Attempt {attempt+1}/12: {e}")
+        time.sleep(10)
+else:
+    raise RuntimeError(
+        f"HF Space at {ENV_URL} is not responding after 2 minutes.\n"
+        "Make sure the Space is Running before submitting this job."
+    )
+
+# ── Install dependencies ───────────────────────────────────────────────────────
+
+print("[HF Job] Installing dependencies...")
+os.system(
+    'pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git" '
+    'trl datasets transformers accelerate scikit-learn pandas numpy matplotlib '
+    '"openenv-core[core]>=0.2.1" --quiet'
+)
+
+# ── Pull latest environment code ──────────────────────────────────────────────
+
+REPO = "https://github.com/CelestialWorthyOfHeavenAndEarth/data-centric-env.git"
+WORK_DIR = "/tmp/data-centric-env"
+
+if not os.path.exists(f"{WORK_DIR}/pyproject.toml"):
+    os.system(f"git clone {REPO} {WORK_DIR}")
+else:
+    os.system(f"git -C {WORK_DIR} pull origin main")
+
+os.chdir(WORK_DIR)
+sys.path.insert(0, WORK_DIR)
+os.system("pip install -e . --quiet")
+print(f"[HF Job] Working dir: {os.getcwd()}")
+os.system("git log --oneline -3")
+
+# ── Set ENV_URL so train_data_centric.py uses the HF Space ───────────────────
+
+os.environ["ENV_URL"] = ENV_URL
+print(f"[HF Job] ENV_URL = {ENV_URL}")
+
+# ── Generate SFT data ─────────────────────────────────────────────────────────
+
+if not os.path.exists("sft_data.jsonl"):
+    print("[HF Job] Generating SFT data...")
+    os.system("python sft_generator.py")
+else:
+    count = sum(1 for _ in open("sft_data.jsonl"))
+    print(f"[HF Job] SFT data exists: {count} examples")
+
+# ── Run full training pipeline ────────────────────────────────────────────────
+
+from train_data_centric import load_model, run_sft_warmup, run_grpo_training, save_model
+import glob
+
+print("[HF Job] Loading model...")
+model, tokenizer = load_model()
+
+print("[HF Job] Phase 1: SFT warmup...")
+model = run_sft_warmup(model, tokenizer)
+print("[HF Job] SFT complete")
+
+print("[HF Job] Phase 2: GRPO training (connecting to HF Space)...")
+resume_from = None
+ckpt_dir = "./data-centric-checkpoints"
+if os.path.exists(ckpt_dir):
+    checkpoints = sorted(glob.glob(f"{ckpt_dir}/checkpoint-*"))
+    if checkpoints:
+        resume_from = checkpoints[-1]
+        print(f"[HF Job] Resuming from: {resume_from}")
+
+model = run_grpo_training(model, tokenizer, resume_from_checkpoint=resume_from)
+print("[HF Job] GRPO complete")
+
+print("[HF Job] Saving model...")
+save_model(model, tokenizer)
+
+print("[HF Job] Generating reward plots...")
+os.system("python plot_rewards.py --log logs/training.jsonl --out plots/")
+
+print("[HF Job] Done! Results in ./logs/ and ./plots/")

inference.py

@@ -0,0 +1,104 @@
+"""
+Heuristic baseline agent for the Data-Centric RL Environment.
+
+Verifies the environment works correctly before any LLM training.
+Run on all 4 tasks, 5 seeds each. Prints a results table.
+"""
+
+import sys
+import os
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from models import DataCentricAction, DataCentricObservation
+from server.data_centric_environment import DataCentricEnvironment
+from server.dataset_generator import TASK_CONFIGS
+
+
+def heuristic_agent(obs: DataCentricObservation, step: int, state: dict) -> str:
+    """
+    Simple heuristic agent that follows:
+    inspect → query_cleaner → apply 1 → apply 2 → validate → query_balancer
+    → apply 1 → validate → submit
+    """
+    if step == 0:
+        return "inspect_dataset"
+    if not state.get("queried_cleaner"):
+        state["queried_cleaner"] = True
+        return "query_cleaner"
+    if state.get("cleaner_applies", 0) < 2:
+        n = state.get("cleaner_applies", 0) + 1
+        state["cleaner_applies"] = n
+        return f"apply {n}"
+    if not state.get("validated"):
+        state["validated"] = True
+        return "validate"
+    if obs.current_accuracy < obs.target_accuracy and not state.get("queried_balancer"):
+        state["queried_balancer"] = True
+        return "query_balancer"
+    if state.get("queried_balancer") and not state.get("balancer_applied"):
+        state["balancer_applied"] = True
+        return "apply 1"
+    if state.get("queried_balancer") and state.get("balancer_applied") and not state.get("validated2"):
+        state["validated2"] = True
+        return "validate"
+    return "submit"
+
+
+def run_heuristic(task: str, seed: int) -> dict:
+    env = DataCentricEnvironment()
+    obs = env.reset(task=task, seed=seed)
+    state = {}
+    total_reward = 0.0
+
+    for step in range(TASK_CONFIGS[task]["budget"]):
+        action_msg = heuristic_agent(obs, step, state)
+        result_obs = env.step(DataCentricAction(message=action_msg))
+        total_reward += result_obs.reward
+        obs = result_obs
+        if obs.done:
+            break
+
+    return {
+        "task": task,
+        "seed": seed,
+        "final_accuracy": obs.current_accuracy,
+        "target": obs.target_accuracy,
+        "hit": obs.current_accuracy >= obs.target_accuracy,
+        "budget_used": obs.step_number,
+        "total_reward": round(total_reward, 4),
+    }
+
+
+def main():
+    tasks = list(TASK_CONFIGS.keys())
+    seeds = [0, 1, 2, 3, 4]
+
+    print(f"\n{'Task':<20} {'Seed':<6} {'Accuracy':<12} {'Target':<10} {'Hit?':<6} {'Budget':<10} {'Reward'}")
+    print("-" * 80)
+
+    hits = 0
+    total = 0
+    for task in tasks:
+        for seed in seeds:
+            r = run_heuristic(task, seed)
+            hit_str = "Y" if r["hit"] else "N"
+            if r["hit"]:
+                hits += 1
+            total += 1
+            print(
+                f"{r['task']:<20} {r['seed']:<6} {r['final_accuracy']:<12.4f} "
+                f"{r['target']:<10.4f} {hit_str:<6} {r['budget_used']:<10} {r['total_reward']}"
+            )
+
+    print("-" * 80)
+    print(f"Hit rate: {hits}/{total} ({100*hits/total:.0f}%)")
+    print()
+    if hits / total >= 0.6:
+        print("  PASS: Heuristic agent validation passed.")
+    else:
+        print("  WARN: Hit rate below 60%. Check environment tuning.")
+
+
+if __name__ == "__main__":
+    main()

models.py

@@ -0,0 +1,104 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Data models for the Data-Centric RL Environment.
+
+Action  → plain text command string (like DataWranglerEnv)
+Observation → rich structured observation with accuracy, quality, budget info
+State   → episode metadata
+"""
+
+from openenv.core.env_server.types import Action, Observation
+from pydantic import Field
+
+
+class DataCentricAction(Action):
+    """Action for the Data-Centric environment — a text command string.
+
+    The agent sends natural-language-style commands to inspect the dataset,
+    query specialist sub-agents, apply their recommendations, and ultimately
+    submit the cleaned dataset for scoring.
+
+    Examples:
+        - "inspect_dataset"
+        - "inspect_model"
+        - "query_cleaner"
+        - "query_augmenter class_1"
+        - "query_balancer"
+        - "query_validator"
+        - "apply 1"
+        - "reject 2"
+        - "validate"
+        - "submit"
+    """
+
+    message: str = Field(..., description="Text command to execute in the environment")
+
+
+class DataCentricObservation(Observation):
+    """Observation returned after each action in the Data-Centric environment.
+
+    Provides the agent with rich feedback about the current episode state,
+    including dataset health, model accuracy, budget, and specialist session info.
+    """
+
+    response: str = Field(
+        default="",
+        description="Text result of the executed command",
+    )
+    current_accuracy: float = Field(
+        default=0.0,
+        description="Last validated model accuracy (or baseline if not yet validated)",
+    )
+    baseline_accuracy: float = Field(
+        default=0.0,
+        description="Accuracy at episode start — never changes",
+    )
+    target_accuracy: float = Field(
+        default=0.0,
+        description="Accuracy threshold the agent must exceed to hit target",
+    )
+    estimated_quality: float = Field(
+        default=0.0,
+        description="Lightweight quality score without sklearn retraining (0.0-1.0)",
+    )
+    dataset_shape: str = Field(
+        default="",
+        description="Current dataset dimensions, e.g. '200 rows × 5 columns'",
+    )
+    rows_preserved_pct: float = Field(
+        default=1.0,
+        description="Fraction of original rows still present (1.0 = no data loss)",
+    )
+    budget_remaining: int = Field(
+        default=0,
+        description="Steps remaining before forced submit",
+    )
+    step_number: int = Field(
+        default=0,
+        description="Current step number in the episode",
+    )
+    max_steps: int = Field(
+        default=30,
+        description="Maximum steps allowed for this task",
+    )
+    active_session: str = Field(
+        default="none",
+        description="Which specialist agent was queried last (cleaner/augmenter/balancer/none)",
+    )
+    validate_calls_remaining: int = Field(
+        default=3,
+        description="How many more free validates remain before reward turns negative",
+    )
+    done: bool = Field(
+        default=False,
+        description="Whether the episode has ended",
+    )
+    reward: float = Field(
+        default=0.0,
+        description="Reward for this step",
+    )

openenv.yaml

@@ -0,0 +1,47 @@
+name: data-centric-env
+description: >
+  Train an LLM agent to improve ML datasets using the Data-Centric AI paradigm.
+  The agent coordinates specialist sub-agents (CleanerAgent, AugmenterAgent,
+  BalancerAgent, ValidatorAgent) to fix data quality issues and improve a fixed
+  classifier's accuracy — without changing the model architecture or hyperparameters.
+  Based on Andrew Ng's Data-Centric AI framework.
+version: "1.1.0"
+tags:
+  - data-centric-ai
+  - world-modeling
+  - professional-tasks
+  - reinforcement-learning
+  - grpo
+  - unsloth
+  - curriculum-learning
+
+tasks:
+  - name: task_0_tutorial
+    description: >
+      Single-issue tutorial. 100 rows, binary classification, 4 features.
+      Only issue: 15-30% missing values. Baseline ~0.62, target 0.73. Budget: 30 steps.
+  - name: task_1_easy
+    description: >
+      Two issues: missing values (8-25%) + mild class imbalance (5-20%).
+      200 rows, binary, 5 features. Baseline ~0.63, target 0.79. Budget: 25 steps.
+  - name: task_2_medium
+    description: >
+      Four issues: missing values, duplicates, class imbalance, type errors.
+      500 rows, 3-class, 7 features. Baseline ~0.58, target 0.74. Budget: 40 steps.
+  - name: task_3_hard
+    description: >
+      Six issues: missing values, duplicates, imbalance, type errors, outliers,
+      cross-column logic errors. 900 rows, 4-class, 10 features.
+      Baseline ~0.54, target 0.71. Budget: 60 steps.
+
+action_type: text
+observation_type: structured
+
+# Reward components:
+#   accuracy_reward:    [-1.0, +1.0] — accuracy delta × 2.5, submit bonus +0.50
+#   process_reward:     [-0.10, +0.05] — workflow pattern scoring
+#   preservation_reward:[-0.40, +0.05] — row preservation incentive
+#   efficiency_reward:  [-0.05, +0.20] — at submit only
+#   step_reward:        [-0.30, +0.15] — proxy quality delta per apply
+# Total clamped to [-1.0, 1.0] by compute_total_reward()
+reward_range: [-1.0, 1.0]

plot_rewards.py

@@ -0,0 +1,259 @@
+"""
+Plot Rewards — Data-Centric AI RL Environment
+=============================================
+Reads JSONL training logs and produces judge-ready plots with labeled axes.
+
+Log format (one JSON object per line in logs/training.jsonl):
+  {
+    "ts": 1714000000.0,       # Unix timestamp
+    "episode": 42,            # Episode number
+    "task": "task_1_easy",    # Task name
+    "level": 1,               # Curriculum level (0=tutorial ... 3=hard)
+    "reward": 0.34,           # Episode reward
+    "accuracy_gain": 0.08,    # Accuracy delta vs baseline
+    "steps_used": 18,         # Steps consumed
+    "success": true           # Reached target accuracy?
+  }
+
+Output (saved to plots/):
+  reward_curve.png     — Episode reward with rolling mean
+  success_rate.png     — Success rate per curriculum level
+  accuracy_gain.png    — Accuracy gain distribution
+  curriculum.png       — Curriculum level over episodes
+
+Usage:
+  python plot_rewards.py                          # default log path
+  python plot_rewards.py --log logs/training.jsonl --out plots/
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import matplotlib
+matplotlib.use("Agg")  # non-interactive backend — safe for headless/Colab
+import matplotlib.pyplot as plt
+import matplotlib.patches as mpatches
+import numpy as np
+import pandas as pd
+
+# ── Style ────────────────────────────────────────────────────────────────────
+
+LEVEL_COLORS = {0: "#4C72B0", 1: "#DD8452", 2: "#55A868", 3: "#C44E52"}
+LEVEL_NAMES  = {0: "tutorial", 1: "easy", 2: "medium", 3: "hard"}
+FIGSIZE = (10, 4)
+DPI = 150
+
+plt.rcParams.update({
+    "font.size": 11,
+    "axes.titlesize": 13,
+    "axes.titleweight": "bold",
+    "axes.labelsize": 11,
+    "grid.alpha": 0.3,
+})
+
+
+# ── Load log ─────────────────────────────────────────────────────────────────
+
+def load_log(log_path: str) -> pd.DataFrame:
+    """Load JSONL training log. Returns empty DataFrame if file not found."""
+    path = Path(log_path)
+    if not path.exists():
+        print(f"[plot_rewards] Log not found: {log_path}")
+        return pd.DataFrame()
+
+    records = []
+    with open(path, encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                try:
+                    records.append(json.loads(line))
+                except json.JSONDecodeError:
+                    pass
+
+    if not records:
+        print(f"[plot_rewards] Log is empty: {log_path}")
+        return pd.DataFrame()
+
+    df = pd.DataFrame(records)
+    # Normalise column names — handle both old and new log formats
+    col_map = {
+        "mean_total_reward": "reward",
+        "mean_env_reward": "accuracy_gain",
+        "stage": "task",
+    }
+    df.rename(columns=col_map, inplace=True)
+    if "episode" not in df.columns:
+        df["episode"] = range(len(df))
+    if "level" not in df.columns:
+        df["level"] = 0
+    if "success" not in df.columns:
+        df["success"] = df.get("accuracy_gain", 0) > 0.05
+    if "accuracy_gain" not in df.columns:
+        df["accuracy_gain"] = 0.0
+    if "reward" not in df.columns:
+        df["reward"] = 0.0
+
+    df.sort_values("episode", inplace=True)
+    df.reset_index(drop=True, inplace=True)
+    print(f"[plot_rewards] Loaded {len(df)} episodes from {log_path}")
+    return df
+
+
+# ── Plots ─────────────────────────────────────────────────────────────────────
+
+def plot_reward_curve(df: pd.DataFrame, out_dir: Path, window: int = 20):
+    """Plot 1: Episode reward over training with rolling mean."""
+    fig, ax = plt.subplots(figsize=FIGSIZE)
+
+    ax.plot(df["episode"], df["reward"], alpha=0.25, color="steelblue",
+            linewidth=0.8, label="Raw reward")
+
+    if len(df) >= window:
+        smooth = df["reward"].rolling(window, min_periods=1).mean()
+        ax.plot(df["episode"], smooth, color="steelblue", linewidth=2.2,
+                label=f"Rolling mean (window={window})")
+
+    # Mark curriculum level transitions
+    level_changes = df[df["level"].diff() != 0]
+    for _, row in level_changes.iterrows():
+        if row["level"] > 0:
+            ax.axvline(row["episode"], color=LEVEL_COLORS.get(int(row["level"]), "gray"),
+                       linewidth=1.0, linestyle="--", alpha=0.7)
+            ax.text(row["episode"] + 0.5, ax.get_ylim()[1] * 0.95,
+                    LEVEL_NAMES.get(int(row["level"]), ""), fontsize=8,
+                    color=LEVEL_COLORS.get(int(row["level"]), "gray"), rotation=90, va="top")
+
+    ax.set_xlabel("Episode")
+    ax.set_ylabel("Episode reward")
+    ax.set_title("Training Reward over Episodes")
+    ax.legend(loc="lower right")
+    ax.grid(True)
+    fig.tight_layout()
+
+    out_path = out_dir / "reward_curve.png"
+    fig.savefig(out_path, dpi=DPI)
+    plt.close(fig)
+    print(f"[plot_rewards] Saved: {out_path}")
+
+
+def plot_success_rate(df: pd.DataFrame, out_dir: Path, window: int = 20):
+    """Plot 2: Success rate per curriculum level."""
+    fig, ax = plt.subplots(figsize=FIGSIZE)
+
+    levels = sorted(df["level"].unique())
+    for level in levels:
+        subset = df[df["level"] == level].copy()
+        subset = subset.sort_values("episode").reset_index(drop=True)
+        rate = subset["success"].rolling(window, min_periods=1).mean()
+        color = LEVEL_COLORS.get(int(level), "gray")
+        label = f"Level {int(level)}: {LEVEL_NAMES.get(int(level), 'unknown')}"
+        ax.plot(subset["episode"], rate, color=color, linewidth=2, label=label)
+
+    ax.axhline(0.60, color="red", linewidth=1.0, linestyle="--", alpha=0.6,
+               label="Advancement threshold (60%)")
+    ax.set_xlabel("Episode")
+    ax.set_ylabel(f"Success rate (rolling mean, window={window})")
+    ax.set_title("Success Rate per Curriculum Level")
+    ax.set_ylim(0, 1.05)
+    ax.legend(loc="lower right")
+    ax.grid(True)
+    fig.tight_layout()
+
+    out_path = out_dir / "success_rate.png"
+    fig.savefig(out_path, dpi=DPI)
+    plt.close(fig)
+    print(f"[plot_rewards] Saved: {out_path}")
+
+
+def plot_accuracy_gain(df: pd.DataFrame, out_dir: Path, window: int = 20):
+    """Plot 3: Accuracy gain over training."""
+    fig, ax = plt.subplots(figsize=FIGSIZE)
+
+    ax.plot(df["episode"], df["accuracy_gain"], alpha=0.25, color="green",
+            linewidth=0.8, label="Raw accuracy gain")
+
+    if len(df) >= window:
+        smooth = df["accuracy_gain"].rolling(window, min_periods=1).mean()
+        ax.plot(df["episode"], smooth, color="green", linewidth=2.2,
+                label=f"Rolling mean (window={window})")
+
+    ax.axhline(0, color="black", linewidth=0.8, linestyle="-", alpha=0.4)
+    ax.set_xlabel("Episode")
+    ax.set_ylabel("Accuracy gain vs baseline")
+    ax.set_title("Accuracy Gain per Episode")
+    ax.legend(loc="lower right")
+    ax.grid(True)
+    fig.tight_layout()
+
+    out_path = out_dir / "accuracy_gain.png"
+    fig.savefig(out_path, dpi=DPI)
+    plt.close(fig)
+    print(f"[plot_rewards] Saved: {out_path}")
+
+
+def plot_curriculum(df: pd.DataFrame, out_dir: Path):
+    """Plot 4: Curriculum level progression over time."""
+    fig, ax = plt.subplots(figsize=FIGSIZE)
+
+    colors = [LEVEL_COLORS.get(int(l), "gray") for l in df["level"]]
+    ax.scatter(df["episode"], df["level"], c=colors, s=4, alpha=0.5, zorder=2)
+
+    # Smooth line
+    ax.plot(df["episode"], df["level"].rolling(10, min_periods=1).mean(),
+            color="black", linewidth=1.5, alpha=0.6, label="Rolling mean level")
+
+    ax.set_xlabel("Episode")
+    ax.set_ylabel("Curriculum level")
+    ax.set_title("Curriculum Progression")
+    ax.set_yticks([0, 1, 2, 3])
+    ax.set_yticklabels(["0: tutorial", "1: easy", "2: medium", "3: hard"])
+    ax.grid(True, axis="x")
+
+    patches = [mpatches.Patch(color=c, label=f"{l}: {LEVEL_NAMES[l]}")
+               for l, c in LEVEL_COLORS.items()]
+    ax.legend(handles=patches, loc="lower right", fontsize=9)
+    fig.tight_layout()
+
+    out_path = out_dir / "curriculum.png"
+    fig.savefig(out_path, dpi=DPI)
+    plt.close(fig)
+    print(f"[plot_rewards] Saved: {out_path}")
+
+
+# ── Entry point ───────────────────────────────────────────────────────────────
+
+def plot_all(log_path: str = "logs/training.jsonl", out_dir: str = "plots/",
+             window: int = 20):
+    df = load_log(log_path)
+    if df.empty:
+        print("[plot_rewards] No data to plot.")
+        return
+
+    out = Path(out_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    plot_reward_curve(df, out, window)
+    plot_success_rate(df, out, window)
+    plot_accuracy_gain(df, out, window)
+    plot_curriculum(df, out)
+
+    print(f"\n[plot_rewards] All plots saved to {out}/")
+    print(f"  Episodes: {len(df)} | "
+          f"Avg reward: {df['reward'].mean():.3f} | "
+          f"Success rate: {df['success'].mean():.1%} | "
+          f"Max level reached: {int(df['level'].max())}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Plot training reward curves")
+    parser.add_argument("--log", default="logs/training.jsonl",
+                        help="Path to JSONL training log")
+    parser.add_argument("--out", default="plots/",
+                        help="Output directory for plots")
+    parser.add_argument("--window", type=int, default=20,
+                        help="Rolling mean window size")
+    args = parser.parse_args()
+    plot_all(args.log, args.out, args.window)

pyproject.toml

@@ -0,0 +1,43 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+[build-system]
+requires = ["setuptools>=45", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "openenv-data_centric_env"
+version = "0.1.0"
+description = "Data Centric Env environment for OpenEnv"
+requires-python = ">=3.10"
+dependencies = [
+    # Core OpenEnv runtime (provides FastAPI server + HTTP client types)
+    "openenv-core[core]>=0.2.2",
+    # Environment-specific dependencies
+    "scikit-learn>=1.3.0",
+    "pandas>=2.0.0",
+    "numpy>=1.24.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+    "pytest-cov>=4.0.0",
+]
+
+[project.scripts]
+# Server entry point - enables running via: uv run --project . server
+# or: python -m data_centric_env.server.app
+server = "data_centric_env.server.app:main"
+
+[tool.setuptools]
+include-package-data = true
+packages = ["data_centric_env", "data_centric_env.server"]
+package-dir = { "data_centric_env" = ".", "data_centric_env.server" = "server" }
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+pythonpath = ["."]

server/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Data-Centric AI RL environment server components."""
+
+from .data_centric_environment import DataCentricEnvironment
+
+__all__ = ["DataCentricEnvironment"]

server/anti_exploit.py

@@ -0,0 +1,202 @@
+"""
+Anti-Exploit Protections for Data-Centric RL Environment.
+
+Centralised module for all anti-hacking checks:
+  1. Input truncation (>200 chars → truncate, -0.02 penalty)
+  2. Validate spam prevention (cooldown + diminishing returns)
+  3. Recommendation ID staleness check
+  4. Ground truth immutability assertion
+  5. Catastrophic data loss detection
+  6. Duplicate apply prevention
+  7. Max applies per session (3)
+  8. Episode wall-clock timeout (5 min → forced submit, -0.10)
+  9. Step timeout (5 sec → timeout obs, -0.05)
+"""
+
+import logging
+import time
+from dataclasses import dataclass, field
+from typing import Optional, Set
+
+logger = logging.getLogger(__name__)
+
+MAX_ACTION_CHARS = 200
+MAX_APPLIES_PER_SESSION = 3
+FREE_VALIDATES = 3
+VALIDATE_COOLDOWN = 2            # must take this many non-validate actions before next validate
+EPISODE_TIMEOUT_SECS = 5 * 60   # 5 minutes
+STEP_TIMEOUT_SECS = 5            # 5 seconds per step
+
+
+# ── Exploit tracker (per episode state) ──────────────────────────────────────
+
+@dataclass
+class AntiExploitState:
+    # Validate tracking
+    validate_call_count: int = 0
+    steps_since_last_validate: int = 0   # cooldown counter
+
+    # Apply tracking
+    applied_ids_this_session: Set[int] = field(default_factory=set)
+    applies_this_session: int = 0
+
+    # Timing
+    episode_start_time: float = field(default_factory=time.time)
+
+    # Ground truth row count (set at reset)
+    ground_truth_row_count: int = 0
+
+
+# ── 1. Input truncation ───────────────────────────────────────────────────────
+
+def check_and_truncate_input(action: str) -> tuple[str, float, bool]:
+    """
+    Returns (truncated_action, penalty, was_truncated).
+    Penalty is -0.02 if truncated, else 0.0.
+    """
+    if len(action) > MAX_ACTION_CHARS:
+        logger.warning(
+            "Input truncated: original length %d > %d", len(action), MAX_ACTION_CHARS
+        )
+        return action[:MAX_ACTION_CHARS], -0.02, True
+    return action, 0.0, False
+
+
+# ── 2. Validate cooldown ──────────────────────────────────────────────────────
+
+def check_validate_cooldown(state: AntiExploitState) -> tuple[bool, str]:
+    """
+    Returns (allowed, error_message).
+    Validate is blocked if steps_since_last_validate < VALIDATE_COOLDOWN.
+    """
+    if state.steps_since_last_validate < VALIDATE_COOLDOWN and state.validate_call_count > 0:
+        return False, (
+            f"Validate on cooldown. Take {VALIDATE_COOLDOWN - state.steps_since_last_validate} "
+            f"more action(s) before validating again."
+        )
+    return True, ""
+
+
+def get_validate_reward(state: AntiExploitState) -> float:
+    """Returns +0.02 for first FREE_VALIDATES calls, -0.01 thereafter."""
+    if state.validate_call_count < FREE_VALIDATES:
+        return 0.02
+    return -0.01
+
+
+def record_validate(state: AntiExploitState):
+    state.validate_call_count += 1
+    state.steps_since_last_validate = 0
+
+
+def record_non_validate_step(state: AntiExploitState):
+    state.steps_since_last_validate += 1
+
+
+# ── 3. Recommendation staleness ───────────────────────────────────────────────
+
+def check_recommendation_staleness(
+    rec_id: int,
+    current_session_id: str,
+    recommendation_session_id: str,
+) -> tuple[bool, str]:
+    """Returns (is_fresh, error_message)."""
+    if current_session_id != recommendation_session_id:
+        return False, (
+            f"Stale recommendation ID {rec_id}. "
+            "Please re-query for fresh recommendations."
+        )
+    return True, ""
+
+
+# ── 4. Ground truth immutability ──────────────────────────────────────────────
+
+def assert_ground_truth_intact(
+    ground_truth_len: int,
+    original_gt_len: int,
+) -> tuple[bool, str]:
+    """Asserts ground truth has not been mutated."""
+    if ground_truth_len != original_gt_len:
+        msg = (
+            f"INTEGRITY VIOLATION: ground_truth row count changed "
+            f"({original_gt_len} → {ground_truth_len}). This should never happen."
+        )
+        logger.critical(msg)
+        return False, msg
+    return True, ""
+
+
+# ── 5. Catastrophic data loss ─────────────────────────────────────────────────
+
+def check_catastrophic_data_loss(
+    current_rows: int,
+    original_rows: int,
+) -> tuple[bool, str]:
+    """Returns (is_catastrophic, message)."""
+    ratio = current_rows / max(original_rows, 1)
+    if ratio < 0.50:
+        msg = (
+            f"CATASTROPHIC DATA LOSS: only {current_rows}/{original_rows} rows remain "
+            f"({ratio*100:.1f}%). Episode terminated."
+        )
+        logger.error(msg)
+        return True, msg
+    return False, ""
+
+
+# ── 6 & 7. Duplicate apply and session limit ──────────────────────────────────
+
+def check_apply_allowed(
+    rec_id: int,
+    state: AntiExploitState,
+) -> tuple[bool, str]:
+    """
+    Returns (allowed, error_message).
+    Blocks: duplicate ID in session, or session apply limit reached.
+    """
+    if state.applies_this_session >= MAX_APPLIES_PER_SESSION:
+        return False, (
+            f"Max {MAX_APPLIES_PER_SESSION} applies per query session reached. "
+            "Please re-query for more options."
+        )
+    if rec_id in state.applied_ids_this_session:
+        return False, (
+            f"Recommendation {rec_id} has already been applied this session. "
+            "Duplicate apply not allowed."
+        )
+    return True, ""
+
+
+def record_apply(rec_id: int, state: AntiExploitState):
+    state.applied_ids_this_session.add(rec_id)
+    state.applies_this_session += 1
+
+
+def reset_session_apply_state(state: AntiExploitState):
+    """Call this whenever a new query_X command resets the session."""
+    state.applied_ids_this_session = set()
+    state.applies_this_session = 0
+
+
+# ── 8. Episode timeout ────────────────────────────────────────────────────────
+
+def check_episode_timeout(state: AntiExploitState) -> tuple[bool, str]:
+    elapsed = time.time() - state.episode_start_time
+    if elapsed > EPISODE_TIMEOUT_SECS:
+        msg = (
+            f"Episode wall-clock timeout ({elapsed:.0f}s > {EPISODE_TIMEOUT_SECS}s). "
+            "Forcing submit. Penalty: -0.10."
+        )
+        logger.warning(msg)
+        return True, msg
+    return False, ""
+
+
+# ── 9. Step timeout context manager ──────────────────────────────────────────
+
+class StepTimeoutError(Exception):
+    pass
+
+
+def validate_calls_remaining(state: AntiExploitState) -> int:
+    return max(0, FREE_VALIDATES - state.validate_call_count)

server/app.py

@@ -0,0 +1,136 @@
+"""FastAPI application for the Data Centric Env Environment."""
+
+import sys
+import os
+
+# Ensure the project root is on the path regardless of how the server is launched
+_HERE = os.path.dirname(os.path.abspath(__file__))
+_ROOT = os.path.dirname(_HERE)
+if _ROOT not in sys.path:
+    sys.path.insert(0, _ROOT)
+
+try:
+    from openenv.core.env_server.http_server import create_app
+except Exception as e:
+    raise ImportError(
+        "openenv-core is required. Install with: pip install openenv-core"
+    ) from e
+
+try:
+    from ..models import DataCentricAction, DataCentricObservation
+    from .data_centric_environment import DataCentricEnvironment
+except (ImportError, ModuleNotFoundError):
+    from models import DataCentricAction, DataCentricObservation
+    from server.data_centric_environment import DataCentricEnvironment
+
+from fastapi.responses import HTMLResponse
+
+# max_concurrent_envs=1: avoids concurrency safety check that instantiates the env
+# at startup (which would load sklearn and pandas, slowing HF health check).
+# Increase if running on a paid Space with more RAM.
+app = create_app(
+    DataCentricEnvironment,
+    DataCentricAction,
+    DataCentricObservation,
+    env_name="data_centric_env",
+    max_concurrent_envs=1,
+)
+
+_LANDING_HTML = """<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="UTF-8">
+<title>Data-Centric AI RL Environment</title>
+<style>
+  body { font-family: system-ui, sans-serif; background: #0f1117; color: #e0e0e0;
+         display: flex; justify-content: center; padding: 60px 20px; margin: 0; }
+  .card { max-width: 700px; width: 100%; }
+  h1 { font-size: 2rem; margin-bottom: 4px; color: #fff; }
+  .badge { display:inline-block; background:#238636; color:#fff; border-radius:12px;
+           padding:2px 10px; font-size:0.8rem; margin-bottom:24px; }
+  p { color: #aaa; line-height: 1.6; }
+  .grid { display: grid; grid-template-columns: 1fr 1fr; gap: 12px; margin: 28px 0; }
+  .endpoint { background: #1c1f26; border: 1px solid #30363d; border-radius: 8px;
+              padding: 14px 18px; }
+  .endpoint code { color: #58a6ff; font-size: 0.9rem; }
+  .endpoint small { color: #666; display:block; margin-top:4px; }
+  a { color: #58a6ff; text-decoration: none; }
+  a:hover { text-decoration: underline; }
+  .footer { margin-top: 32px; font-size: 0.8rem; color: #555; border-top: 1px solid #21262d; padding-top: 16px; }
+</style>
+</head>
+<body>
+<div class="card">
+  <h1>&#x1F9E0; Data-Centric AI Environment</h1>
+  <span class="badge">&#x2022; Running</span>
+  <p>An <a href="https://github.com/meta-pytorch/OpenEnv" target="_blank">OpenEnv</a>-compliant
+  RL environment that trains an LLM to coordinate specialist data-cleaning agents
+  to improve a frozen ML classifier — without touching the model.</p>
+
+  <div class="grid">
+    <div class="endpoint">
+      <code>GET /health</code>
+      <small>Server health check</small>
+    </div>
+    <div class="endpoint">
+      <code>GET /docs</code>
+      <small>Interactive API docs (Swagger)</small>
+    </div>
+    <div class="endpoint">
+      <code>POST /reset</code>
+      <small>Start a new episode</small>
+    </div>
+    <div class="endpoint">
+      <code>POST /step</code>
+      <small>Execute an action</small>
+    </div>
+    <div class="endpoint">
+      <code>WS /ws</code>
+      <small>WebSocket (stateful session)</small>
+    </div>
+    <div class="endpoint">
+      <code>GET /state</code>
+      <small>Current episode state</small>
+    </div>
+  </div>
+
+  <p><strong>Quick start:</strong></p>
+  <pre style="background:#161b22;padding:14px;border-radius:8px;font-size:0.85rem;overflow:auto">
+pip install git+https://huggingface.co/spaces/Aswini-Kumar/data-centric-env
+
+from data_centric_env import DataCentricEnv, DataCentricAction
+with DataCentricEnv(base_url="https://aswini-kumar-data-centric-env.hf.space").sync() as env:
+    obs = env.reset(task="task_0_tutorial", seed=42)
+    result = env.step(DataCentricAction(message="query_analyst"))
+    print(result.observation.response)</pre>
+
+  <div class="footer">
+    <a href="/docs">API Docs</a> &nbsp;|&nbsp;
+    <a href="/health">Health</a> &nbsp;|&nbsp;
+    <a href="https://github.com/CelestialWorthyOfHeavenAndEarth/data-centric-env" target="_blank">GitHub</a> &nbsp;|&nbsp;
+    <a href="https://huggingface.co/spaces/Aswini-Kumar/data-centric-env" target="_blank">HF Space</a>
+  </div>
+</div>
+</body>
+</html>"""
+
+
+@app.get("/", response_class=HTMLResponse, include_in_schema=False)
+@app.get("/web", response_class=HTMLResponse, include_in_schema=False)
+async def landing():
+    """Human-readable landing page for the HF Space App tab."""
+    return HTMLResponse(content=_LANDING_HTML)
+
+
+def main(host: str = "0.0.0.0", port: int = 7860):
+    import uvicorn
+    uvicorn.run(app, host=host, port=port)
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", default="0.0.0.0")
+    parser.add_argument("--port", type=int, default=7860)
+    args = parser.parse_args()
+    main(host=args.host, port=args.port)

server/data_centric_environment.py

@@ -0,0 +1,727 @@
+"""Main DataCentric RL Environment."""
+
+import logging
+import time
+from copy import deepcopy
+from typing import Any, Dict, List, Optional
+from uuid import uuid4
+
+import pandas as pd
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import State
+
+try:
+    from ..models import DataCentricAction, DataCentricObservation
+except ImportError:
+    try:
+        from models import DataCentricAction, DataCentricObservation
+    except ImportError:
+        import sys as _sys, os as _os
+        _sys.path.insert(0, _os.path.dirname(_os.path.dirname(_os.path.abspath(__file__))))
+        from models import DataCentricAction, DataCentricObservation
+
+try:
+    from .anti_exploit import (
+        AntiExploitState, assert_ground_truth_intact,
+        check_and_truncate_input, check_apply_allowed,
+        check_catastrophic_data_loss, check_episode_timeout,
+        check_validate_cooldown, get_validate_reward, record_apply,
+        record_non_validate_step, record_validate, reset_session_apply_state,
+        validate_calls_remaining,
+    )
+    from .dataset_generator import TASK_CONFIGS, generate_dataset
+    from .grader import (
+        compute_accuracy_reward, compute_efficiency_reward,
+        compute_lightweight_score, compute_preservation_reward,
+        compute_process_reward, compute_step_reward, compute_total_reward,
+    )
+    from .model_evaluator import ModelEvaluator
+    from .specialist_agents import (
+        AugmenterAgent, AnalystAgent, BalancerAgent, CleanerAgent,
+        SessionRegistry, ValidatorAgent, compute_drift, format_drift_summary,
+    )
+except ImportError:
+    from server.anti_exploit import (
+        AntiExploitState, assert_ground_truth_intact,
+        check_and_truncate_input, check_apply_allowed,
+        check_catastrophic_data_loss, check_episode_timeout,
+        check_validate_cooldown, get_validate_reward, record_apply,
+        record_non_validate_step, record_validate, reset_session_apply_state,
+        validate_calls_remaining,
+    )
+    from server.dataset_generator import TASK_CONFIGS, generate_dataset
+    from server.grader import (
+        compute_accuracy_reward, compute_efficiency_reward,
+        compute_lightweight_score, compute_preservation_reward,
+        compute_process_reward, compute_step_reward, compute_total_reward,
+    )
+    from server.model_evaluator import ModelEvaluator
+    from server.specialist_agents import (
+        AugmenterAgent, AnalystAgent, BalancerAgent, CleanerAgent,
+        SessionRegistry, ValidatorAgent, compute_drift, format_drift_summary,
+    )
+
+logger = logging.getLogger(__name__)
+
+AVAILABLE_COMMANDS = """Available commands:
+  inspect_dataset          — shape, dtypes, missing, class distribution
+  inspect_model            — accuracy (RF + LR), F1, feature importance
+  query_analyst            — holistic diagnosis + prioritised action plan (costs 1 budget)
+  query_cleaner            — get cleaning recommendations
+  query_augmenter [class]  — get augmentation suggestions
+  query_balancer           — get resampling recommendations
+  query_validator          — check rule violations (costs 2 budget)
+  apply [id]               — apply recommendation by ID
+  reject [id]              — reject a recommendation
+  undo                     — revert last apply (max 3 levels)
+  validate                 — retrain and score (cooldown applies)
+  submit                   — finalize episode"""
+
+
+class DataCentricEnvironment(Environment):
+    """Data-Centric AI RL Environment."""
+
+    SUPPORTS_CONCURRENT_SESSIONS: bool = True
+
+    def __init__(self):
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        self._ground_truth: Optional[pd.DataFrame] = None
+        self._working_copy: Optional[pd.DataFrame] = None
+        self._metadata: Dict[str, Any] = {}
+        self._action_history: List[str] = []
+        self._exploit: Optional[AntiExploitState] = None
+        # fast_mode=True: uses n_estimators=20 for training rollouts (~4x faster)
+        self._evaluator = ModelEvaluator(fast_mode=True)
+        self._session_registry = SessionRegistry()
+        self._cleaner = CleanerAgent()
+        self._augmenter = AugmenterAgent()
+        self._balancer = BalancerAgent()
+        self._validator = ValidatorAgent()
+        self._analyst = AnalystAgent()
+        self._current_accuracy: float = 0.0
+        self._previous_accuracy: float = 0.0
+        self._active_session: str = "none"
+        self._task: str = "task_0_tutorial"
+        # Snapshot stack for undo command (max 3 snapshots)
+        self._dataset_history: List[pd.DataFrame] = []
+        self._max_history: int = 3
+
+    # ── reset ────────────────────────────────────────────────────────────────
+
+    def reset(self, task: str = "task_0_tutorial", seed: int = 42) -> DataCentricObservation:
+        self._task = task if task in TASK_CONFIGS else "task_0_tutorial"
+        cfg = TASK_CONFIGS[self._task]
+
+        self._ground_truth, self._working_copy, self._metadata = generate_dataset(
+            self._task, seed=seed
+        )
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        self._action_history = []
+        self._exploit = AntiExploitState(
+            episode_start_time=time.time(),
+            ground_truth_row_count=len(self._ground_truth),
+        )
+        self._evaluator.invalidate_cache()
+        self._session_registry = SessionRegistry()
+        self._active_session = "none"
+        self._dataset_history = []   # clear snapshot stack on reset
+        reset_session_apply_state(self._exploit)
+
+        # Store episode-start missing count for quality score baseline
+        self._metadata["initial_missing"] = int(self._working_copy.isnull().sum().sum())
+        self._metadata["baseline_accuracy"] = cfg["baseline_accuracy"]
+
+        baseline = cfg["baseline_accuracy"]
+        self._current_accuracy = baseline
+        self._previous_accuracy = baseline
+        quality = compute_lightweight_score(
+            self._working_copy, self._ground_truth,
+            self._metadata["original_length"], self._metadata["col_meta"],
+            initial_missing=self._metadata["initial_missing"],
+        )
+        wc = self._working_copy
+        return DataCentricObservation(
+            response=(
+                f"Episode started: {self._task}\n"
+                f"Baseline accuracy: {baseline:.4f} | Target: {cfg['target_accuracy']:.4f}\n"
+                f"Dataset: {len(wc)} rows × {len(wc.columns)-1} features\n"
+                f"Budget: {cfg['budget']} steps\n\n{AVAILABLE_COMMANDS}"
+            ),
+            current_accuracy=baseline,
+            baseline_accuracy=baseline,
+            target_accuracy=cfg["target_accuracy"],
+            estimated_quality=quality,
+            dataset_shape=f"{len(wc)} rows × {len(wc.columns)-1} columns",
+            rows_preserved_pct=1.0,
+            budget_remaining=cfg["budget"],
+            step_number=0,
+            max_steps=cfg["budget"],
+            active_session="none",
+            validate_calls_remaining=3,
+            done=False,
+            reward=0.0,
+        )
+
+    # ── step ─────────────────────────────────────────────────────────────────
+
+    def step(self, action: DataCentricAction) -> DataCentricObservation:
+        if self._working_copy is None:
+            return self._error_obs("Call reset() first.")
+
+        # Episode timeout
+        timeout, tmsg = check_episode_timeout(self._exploit)
+        if timeout:
+            return self._do_submit(penalty=-0.10, extra_msg=tmsg)
+
+        # Input truncation
+        raw_msg = action.message
+        msg, trunc_penalty, was_truncated = check_and_truncate_input(raw_msg)
+        if was_truncated:
+            logger.warning("Input truncated.")
+
+        cfg = TASK_CONFIGS[self._task]
+        self._state.step_count += 1
+        step_num = self._state.step_count
+        budget_remaining = cfg["budget"] - step_num
+        cmd_parts = msg.strip().split()
+        cmd = cmd_parts[0].lower() if cmd_parts else ""
+
+        # Out of budget → force submit
+        if budget_remaining < 0:
+            return self._do_submit(penalty=0.0, extra_msg="Budget exhausted.")
+
+        # Record action
+        self._action_history.append(msg)
+
+        # Process reward component (computed for all actions)
+        r_process = compute_process_reward(self._action_history[:-1], msg)
+
+        # Route command
+        if cmd == "inspect_dataset":
+            obs = self._cmd_inspect_dataset(step_num, budget_remaining, r_process, trunc_penalty)
+        elif cmd == "inspect_model":
+            obs = self._cmd_inspect_model(step_num, budget_remaining, r_process, trunc_penalty)
+        elif cmd == "query_cleaner":
+            obs = self._cmd_query_cleaner(step_num, budget_remaining, r_process, trunc_penalty)
+        elif cmd == "query_augmenter":
+            cls = cmd_parts[1] if len(cmd_parts) > 1 else None
+            obs = self._cmd_query_augmenter(cls, step_num, budget_remaining, r_process, trunc_penalty)
+        elif cmd == "query_balancer":
+            obs = self._cmd_query_balancer(step_num, budget_remaining, r_process, trunc_penalty)
+        elif cmd == "query_analyst":
+            obs = self._cmd_query_analyst(step_num, budget_remaining, r_process, trunc_penalty)
+        elif cmd == "query_validator":
+            obs = self._cmd_query_validator(step_num, budget_remaining, r_process, trunc_penalty)
+        elif cmd == "apply":
+            try:
+                rec_id = int(cmd_parts[1]) if len(cmd_parts) > 1 else -1
+            except ValueError:
+                rec_id = -1
+            obs = self._cmd_apply(rec_id, step_num, budget_remaining, r_process, trunc_penalty)
+        elif cmd == "reject":
+            try:
+                rec_id = int(cmd_parts[1]) if len(cmd_parts) > 1 else -1
+            except ValueError:
+                rec_id = -1
+            obs = self._cmd_reject(rec_id, step_num, budget_remaining, r_process, trunc_penalty)
+        elif cmd == "validate":
+            obs = self._cmd_validate(step_num, budget_remaining, r_process, trunc_penalty)
+        elif cmd == "submit":
+            obs = self._do_submit()
+        elif cmd == "undo":
+            obs = self._cmd_undo(step_num, budget_remaining, r_process, trunc_penalty)
+        else:
+            obs = self._unknown_cmd_obs(msg, step_num, budget_remaining, r_process + trunc_penalty)
+
+        if cmd != "validate":
+            record_non_validate_step(self._exploit)
+
+        return obs
+
+    # ── command handlers ─────────────────────────────────────────────────────
+
+    def _cmd_inspect_dataset(self, step, budget, r_process, trunc_pen) -> DataCentricObservation:
+        wc = self._working_copy
+        orig_len = self._metadata["original_length"]
+        missing = wc.isnull().sum()
+        missing_str = "\n".join(f"  {c}: {v}" for c, v in missing.items() if v > 0) or "  None"
+        vc = wc["target"].value_counts().sort_index()
+        class_str = ", ".join(f"class {k}: {v}" for k, v in vc.items())
+        rows_pct = len(wc) / orig_len
+        response = (
+            f"=== Dataset Inspection ===\n"
+            f"Shape: {len(wc)} rows × {len(wc.columns)-1} features\n"
+            f"Original rows: {orig_len} | Preserved: {rows_pct*100:.1f}%\n"
+            f"Duplicates: {wc.duplicated().sum()}\n"
+            f"Missing values:\n{missing_str}\n"
+            f"Class distribution: {class_str}\n"
+            f"Dtypes: {dict(wc.dtypes.astype(str))}"
+        )
+        reward = compute_total_reward(0.0, r_process, 0.0) + trunc_pen
+        return self._make_obs(response, step, budget, reward)
+
+    def _cmd_inspect_model(self, step, budget, r_process, trunc_pen) -> DataCentricObservation:
+        acc, per_class, from_cache, lr_acc = self._evaluator.evaluate(
+            self._working_copy, self._ground_truth
+        )
+        cache_label = " (cached)" if from_cache else ""
+        lines = [f"=== Model Inspection{cache_label} ===",
+                 f"RF Accuracy:  {acc:.4f}",
+                 f"LR Accuracy:  {lr_acc:.4f}  (secondary — diagnostic only)"]
+        for cls, metrics in per_class.items():
+            if isinstance(metrics, dict):
+                lines.append(
+                    f"  Class {cls}: precision={metrics.get('precision',0):.3f} "
+                    f"recall={metrics.get('recall',0):.3f} "
+                    f"f1={metrics.get('f1-score',0):.3f}"
+                )
+        feat_text = self._evaluator.feature_importance_text()
+        if feat_text:
+            lines.append(feat_text)
+        response = "\n".join(lines)
+        reward = compute_total_reward(0.0, r_process, 0.0) + trunc_pen
+        return self._make_obs(response, step, budget, reward)
+
+    def _cmd_query_cleaner(self, step, budget, r_process, trunc_pen) -> DataCentricObservation:
+        reset_session_apply_state(self._exploit)
+        recs = self._cleaner.query(
+            self._working_copy, self._session_registry, self._metadata["col_meta"]
+        )
+        self._active_session = f"cleaner:{self._session_registry.current_session_id[:8]}"
+        lines = ["=== Cleaner Recommendations ==="]
+        for r in recs:
+            lines.append(
+                f"[{r.id}] {r.description}\n"
+                f"    type={r.action_type} impact={r.estimated_impact:+.3f} "
+                f"confidence={r.confidence:.2f}"
+            )
+        response = "\n".join(lines) if recs else "No cleaning issues detected."
+        reward = compute_total_reward(0.0, r_process, 0.0) + trunc_pen
+        return self._make_obs(response, step, budget, reward)
+
+    def _cmd_query_augmenter(self, cls, step, budget, r_process, trunc_pen) -> DataCentricObservation:
+        reset_session_apply_state(self._exploit)
+        recs = self._augmenter.query(self._working_copy, self._session_registry, cls)
+        self._active_session = f"augmenter:{self._session_registry.current_session_id[:8]}"
+        lines = ["=== Augmenter Recommendations ==="]
+        for r in recs:
+            lines.append(
+                f"[{r.id}] {r.description}\n"
+                f"    type={r.action_type} impact={r.estimated_impact:+.3f} "
+                f"confidence={r.confidence:.2f}"
+            )
+        response = "\n".join(lines) if recs else "No augmentation needed."
+        reward = compute_total_reward(0.0, r_process, 0.0) + trunc_pen
+        return self._make_obs(response, step, budget, reward)
+
+    def _cmd_query_balancer(self, step, budget, r_process, trunc_pen) -> DataCentricObservation:
+        reset_session_apply_state(self._exploit)
+        recs = self._balancer.query(self._working_copy, self._session_registry)
+        self._active_session = f"balancer:{self._session_registry.current_session_id[:8]}"
+        lines = ["=== Balancer Recommendations ==="]
+        for r in recs:
+            lines.append(
+                f"[{r.id}] {r.description}\n"
+                f"    type={r.action_type} impact={r.estimated_impact:+.3f} "
+                f"confidence={r.confidence:.2f}"
+            )
+        response = "\n".join(lines) if recs else "Dataset is already balanced."
+        reward = compute_total_reward(0.0, r_process, 0.0) + trunc_pen
+        return self._make_obs(response, step, budget, reward)
+
+    def _cmd_query_analyst(self, step, budget, r_process, trunc_pen) -> DataCentricObservation:
+        """Holistic diagnosis + prioritised action plan. Costs 1 budget."""
+        # Costs 1 extra budget step
+        self._state.step_count += 1
+        plan = self._analyst.query(
+            self._working_copy,
+            self._metadata["col_meta"],
+            self._current_accuracy,
+            TASK_CONFIGS[self._task]["target_accuracy"],
+            budget - 1,
+        )
+        response = f"=== Analyst Report (costs 1 budget) ===\n{plan}"
+        reward = compute_total_reward(0.0, r_process + 0.02, 0.0) + trunc_pen  # small bonus for planning
+        budget_remaining = TASK_CONFIGS[self._task]["budget"] - self._state.step_count
+        return self._make_obs(response, step, budget_remaining, reward)
+
+    def _cmd_query_validator(self, step, budget, r_process, trunc_pen) -> DataCentricObservation:
+        # Costs 2 budget
+        self._state.step_count += 1
+        violations = self._validator.query(self._working_copy, self._metadata["col_meta"])
+        lines = ["=== Validator Report (costs 2 budget) ==="]
+        if violations:
+            for v in violations:
+                lines.append(
+                    f"  [{v.severity}] [{v.column}] rule={v.rule} count={v.count}\n    {v.description}"
+                )
+        else:
+            lines.append("  No rule violations found.")
+        response = "\n".join(lines)
+        reward = compute_total_reward(0.0, r_process, 0.0) + trunc_pen
+        budget_remaining = TASK_CONFIGS[self._task]["budget"] - self._state.step_count
+        return self._make_obs(response, step, budget_remaining, reward)
+
+    def _cmd_apply(self, rec_id, step, budget, r_process, trunc_pen) -> DataCentricObservation:
+        if rec_id < 1:
+            # Error: return 0 reward (no penalty, no bonus)
+            return self._make_obs("Error: invalid recommendation ID.", step, budget, 0.0)
+
+        # Check apply allowed (duplicate / session limit) — 0 reward on error
+        allowed, err = check_apply_allowed(rec_id, self._exploit)
+        if not allowed:
+            return self._make_obs(f"Error: {err}", step, budget, 0.0)
+
+        # Get recommendation (staleness check) — 0 reward, no penalty
+        rec = self._session_registry.get(rec_id, self._session_registry.current_session_id)
+        if rec is None:
+            return self._make_obs(
+                f"Error: stale recommendation ID {rec_id}. Please re-query for fresh recommendations.",
+                step, budget, 0.0
+            )
+
+        # Capture quality before mutation for step reward
+        quality_before = compute_lightweight_score(
+            self._working_copy, self._ground_truth,
+            self._metadata["original_length"], self._metadata["col_meta"],
+            initial_missing=self._metadata.get("initial_missing"),
+        )
+
+        # Execute payload
+        payload = rec._payload
+        action_type = payload.get("action", "")
+        wc = self._working_copy
+        orig_len = self._metadata["original_length"]
+        pre_rows = len(wc)
+        pre_missing = int(wc.isnull().sum().sum())
+        pre_dups = int(wc.duplicated().sum())
+
+        # Save snapshot for undo before mutating
+        self._dataset_history.append(self._working_copy.copy())
+        if len(self._dataset_history) > self._max_history:
+            self._dataset_history.pop(0)
+
+        try:
+            if action_type == "fill_missing":
+                col = payload["column"]
+                strategy = payload.get("strategy", "mean")  # honor smarter CleanerAgent choice
+                numeric = pd.to_numeric(wc[col], errors="coerce")
+                if strategy == "median":
+                    fill_val = float(numeric.median())
+                else:
+                    fill_val = float(numeric.mean())
+                wc[col] = numeric.fillna(fill_val)
+                self._working_copy = wc
+
+            elif action_type == "remove_duplicates":
+                self._working_copy = wc.drop_duplicates().reset_index(drop=True)
+
+            elif action_type == "fix_type_errors":
+                col = payload["column"]
+                numeric = pd.to_numeric(wc[col], errors="coerce")
+                mean_val = float(numeric.mean())
+                wc[col] = numeric.fillna(mean_val)
+                self._working_copy = wc
+
+            elif action_type == "augment_class":
+                cls_int = payload["class"]
+                n_synth = payload["n_synth"]
+                cls_rows = wc[wc["target"] == cls_int]
+                if len(cls_rows) > 0:
+                    synth = cls_rows.sample(n=n_synth, replace=True, random_state=42)
+                    noise_cols = [c for c in synth.columns if c != "target"]
+                    for c in noise_cols:
+                        try:
+                            synth[c] = pd.to_numeric(synth[c], errors="coerce")
+                            synth[c] = synth[c] + synth[c].std() * 0.1
+                        except Exception:
+                            pass
+                    self._working_copy = pd.concat([wc, synth], ignore_index=True)
+
+            elif action_type == "oversample":
+                cls_int = payload["class"]
+                target_count = payload["target_count"]
+                cls_rows = wc[wc["target"] == cls_int]
+                n_needed = max(0, target_count - len(cls_rows))
+                if n_needed > 0:
+                    extra = cls_rows.sample(n=n_needed, replace=True, random_state=42)
+                    self._working_copy = pd.concat([wc, extra], ignore_index=True)
+
+            elif action_type == "undersample":
+                cls_int = payload["class"]
+                target_count = payload["target_count"]
+                cls_rows = wc[wc["target"] == cls_int]
+                if len(cls_rows) > target_count:
+                    keep = cls_rows.sample(n=target_count, random_state=42)
+                    other = wc[wc["target"] != cls_int]
+                    self._working_copy = pd.concat([keep, other], ignore_index=True)
+
+            elif action_type == "remove_outlier_rows":
+                col = payload["column"]
+                pct = payload.get("pct", 5)
+                try:
+                    numeric = pd.to_numeric(wc[col], errors="coerce")
+                    threshold = float(numeric.quantile(pct / 100))
+                    self._working_copy = wc[pd.to_numeric(wc[col], errors="coerce") >= threshold].reset_index(drop=True)
+                except Exception:
+                    pass
+
+        except Exception as exc:
+            logger.exception("Error executing apply: %s", exc)
+            return self._make_obs(f"Error executing recommendation: {exc}", step, budget, 0.0)
+
+        record_apply(rec_id, self._exploit)
+
+        # Ground truth immutability assertion — must never change
+        gt_ok, gt_msg = assert_ground_truth_intact(
+            len(self._ground_truth), self._exploit.ground_truth_row_count
+        )
+        if not gt_ok:
+            logger.critical(gt_msg)
+            return self._do_submit(penalty=-1.0, extra_msg=gt_msg)
+
+        wc_new = self._working_copy
+        post_rows = len(wc_new)
+        post_missing = int(wc_new.isnull().sum().sum())
+        post_dups = int(wc_new.duplicated().sum())
+        rows_pct = post_rows / orig_len
+
+        # Catastrophic data loss
+        catastro, cmsg = check_catastrophic_data_loss(post_rows, orig_len)
+        if catastro:
+            return self._do_submit(penalty=-0.40, extra_msg=cmsg)
+
+        # Preservation reward
+        r_preservation = compute_preservation_reward(post_rows, orig_len)
+
+        # Lightweight quality (use episode-start missing count as denominator)
+        quality = compute_lightweight_score(
+            wc_new, self._ground_truth, orig_len, self._metadata["col_meta"],
+            initial_missing=self._metadata.get("initial_missing"),
+        )
+
+        # Build rich feedback with drift detection
+        cfg = TASK_CONFIGS[self._task]
+        missing_status = "OK" if post_missing == 0 else f"{post_missing} remaining"
+        dup_status = "OK" if post_dups == 0 else f"{post_dups} remaining"
+        drift = compute_drift(self._working_copy, self._ground_truth)
+        drift_summary = format_drift_summary(drift)
+        response = (
+            f"Applied: {action_type} [{rec.description[:80]}]\n\n"
+            f"Dataset health check:\n"
+            f"  Missing values: {missing_status} (was {pre_missing})\n"
+            f"  Duplicates: {dup_status} (was {pre_dups})\n"
+            f"  Row count: {post_rows}/{orig_len} ({rows_pct*100:.1f}% preserved)\n"
+            f"  {drift_summary}\n\n"
+            f"Estimated quality score: {quality:.4f}\n"
+            f"Budget remaining: {budget}"
+        )
+
+        reward = compute_total_reward(
+            0.0, r_process, r_preservation,
+            reward_step=compute_step_reward(
+                f"apply {rec_id}", quality_before, quality, rows_pct
+            ),
+        ) + trunc_pen
+        self._evaluator.invalidate_cache()
+        return self._make_obs(response, step, budget, reward, quality=quality,
+                              rows_pct=rows_pct)
+
+    def _cmd_reject(self, rec_id, step, budget, r_process, trunc_pen) -> DataCentricObservation:
+        response = (
+            f"Recommendation {rec_id} rejected. It will not appear in future queries."
+            if rec_id >= 1 else "Error: invalid recommendation ID."
+        )
+        reward = compute_total_reward(0.0, r_process + 0.01, 0.0) + trunc_pen
+        return self._make_obs(response, step, budget, reward)
+
+    def _cmd_undo(self, step, budget, r_process, trunc_pen) -> DataCentricObservation:
+        """Restore previous dataset state (max 3 levels deep)."""
+        if self._dataset_history:
+            self._working_copy = self._dataset_history.pop()
+            self._evaluator.invalidate_cache()
+            orig_len = self._metadata["original_length"]
+            rows_pct = len(self._working_copy) / orig_len
+            quality = compute_lightweight_score(
+                self._working_copy, self._ground_truth,
+                orig_len, self._metadata["col_meta"],
+                initial_missing=self._metadata.get("initial_missing"),
+            )
+            response = (
+                f"Undo successful. Reverted to previous dataset state.\n"
+                f"Row count: {len(self._working_copy)}/{orig_len} ({rows_pct*100:.1f}% preserved)\n"
+                f"Estimated quality: {quality:.4f}\n"
+                f"Snapshots remaining: {len(self._dataset_history)}"
+            )
+            reward = compute_total_reward(0.0, r_process - 0.03, 0.0) + trunc_pen  # small cost
+        else:
+            response = "Nothing to undo. No previous state available."
+            reward = compute_total_reward(0.0, r_process - 0.05, 0.0) + trunc_pen  # larger cost
+        return self._make_obs(response, step, budget, reward)
+
+
+    def _cmd_validate(self, step, budget, r_process, trunc_pen) -> DataCentricObservation:
+        allowed, cooldown_msg = check_validate_cooldown(self._exploit)
+        if not allowed:
+            return self._make_obs(cooldown_msg, step, budget, 0.0)
+
+        prev_rf = self._evaluator.last_accuracy
+        prev_lr = self._evaluator.last_lr_accuracy
+
+        acc, per_class, from_cache, lr_acc = self._evaluator.evaluate(
+            self._working_copy, self._ground_truth
+        )
+        cache_label = " (cached)" if from_cache else ""
+
+        if from_cache:
+            r_validate = 0.0
+        else:
+            r_validate = get_validate_reward(self._exploit)
+            record_validate(self._exploit)
+
+        r_accuracy = compute_accuracy_reward(
+            acc, self._current_accuracy,
+            self._metadata["baseline_accuracy"],
+            TASK_CONFIGS[self._task]["target_accuracy"],
+        )
+        self._previous_accuracy = self._current_accuracy
+        self._current_accuracy = acc
+
+        target = TASK_CONFIGS[self._task]["target_accuracy"]
+        agreement = self._evaluator.agreement_signal(acc, lr_acc, prev_rf, prev_lr)
+        feat_text = self._evaluator.feature_importance_text()
+
+        lines = [
+            f"=== Validate{cache_label} ===",
+            f"RF Accuracy: {acc:.4f}  (primary)",
+            f"LR Accuracy: {lr_acc:.4f}  (secondary)",
+            f"Agreement:   {agreement}",
+        ]
+        for cls, metrics in per_class.items():
+            if isinstance(metrics, dict):
+                lines.append(
+                    f"  Class {cls}: p={metrics.get('precision',0):.3f} "
+                    f"r={metrics.get('recall',0):.3f} f1={metrics.get('f1-score',0):.3f}"
+                )
+        lines.append(f"Target: {target:.4f} | {'HIT ✓' if acc >= target else 'Not yet'}")
+        if feat_text:
+            lines.append(feat_text)
+        response = "\n".join(lines)
+
+        reward = compute_total_reward(r_accuracy, r_process + r_validate, 0.0) + trunc_pen
+        return self._make_obs(response, step, budget, reward)
+
+    # ── submit ────────────────────────────────────────────────────────────────
+
+    def _do_submit(self, penalty: float = 0.0, extra_msg: str = "") -> DataCentricObservation:
+        cfg = TASK_CONFIGS[self._task]
+        orig_len = self._metadata["original_length"]
+        budget_remaining = cfg["budget"] - self._state.step_count
+
+        # Final accuracy
+        acc, per_class, _, lr_acc = self._evaluator.evaluate(
+            self._working_copy, self._ground_truth
+        )
+        self._current_accuracy = acc
+
+        r_accuracy = compute_accuracy_reward(
+            acc, self._previous_accuracy,
+            cfg["baseline_accuracy"], cfg["target_accuracy"],
+            is_submit=True,
+        )
+        r_process = compute_process_reward(self._action_history[:-1], "submit")
+        r_preservation = compute_preservation_reward(len(self._working_copy), orig_len)
+        r_efficiency = compute_efficiency_reward(
+            acc, cfg["baseline_accuracy"], cfg["budget"], max(budget_remaining, 0)
+        )
+
+        total = compute_total_reward(r_accuracy, r_process, r_preservation, r_efficiency)
+        total += penalty
+
+        hit = acc >= cfg["target_accuracy"]
+        response = (
+            f"{'=' * 40}\n"
+            f"EPISODE COMPLETE\n"
+            f"{'=' * 40}\n"
+            f"Final accuracy:  {acc:.4f}\n"
+            f"Target accuracy: {cfg['target_accuracy']:.4f}\n"
+            f"Baseline:        {cfg['baseline_accuracy']:.4f}\n"
+            f"Result: {'TARGET HIT ✓' if hit else 'Target not reached'}\n\n"
+            f"Reward breakdown:\n"
+            f"  Accuracy:     {r_accuracy:+.4f}\n"
+            f"  Process:      {r_process:+.4f}\n"
+            f"  Preservation: {r_preservation:+.4f}\n"
+            f"  Efficiency:   {r_efficiency:+.4f}\n"
+            f"  Penalty:      {penalty:+.4f}\n"
+            f"  TOTAL:        {total:+.4f}\n"
+            + (f"\n{extra_msg}" if extra_msg else "")
+        )
+
+        quality = compute_lightweight_score(
+            self._working_copy, self._ground_truth,
+            orig_len, self._metadata["col_meta"],
+        )
+        rows_pct = len(self._working_copy) / orig_len
+
+        return DataCentricObservation(
+            response=response,
+            current_accuracy=acc,
+            baseline_accuracy=cfg["baseline_accuracy"],
+            target_accuracy=cfg["target_accuracy"],
+            estimated_quality=quality,
+            dataset_shape=f"{len(self._working_copy)} rows × {len(self._working_copy.columns)-1} columns",
+            rows_preserved_pct=rows_pct,
+            budget_remaining=max(budget_remaining, 0),
+            step_number=self._state.step_count,
+            max_steps=cfg["budget"],
+            active_session=self._active_session,
+            validate_calls_remaining=validate_calls_remaining(self._exploit),
+            done=True,
+            reward=round(total, 4),
+        )
+
+    # ── helpers ───────────────────────────────────────────────────────────────
+
+    def _make_obs(self, response: str, step: int, budget: int, reward: float,
+                  quality: Optional[float] = None, rows_pct: Optional[float] = None
+                  ) -> DataCentricObservation:
+        cfg = TASK_CONFIGS[self._task]
+        orig_len = self._metadata["original_length"]
+        wc = self._working_copy
+        if quality is None:
+            quality = compute_lightweight_score(
+                wc, self._ground_truth, orig_len, self._metadata["col_meta"],
+                initial_missing=self._metadata.get("initial_missing"),
+            )
+        if rows_pct is None:
+            rows_pct = len(wc) / orig_len
+
+        return DataCentricObservation(
+            response=response,
+            current_accuracy=self._current_accuracy,
+            baseline_accuracy=cfg["baseline_accuracy"],
+            target_accuracy=cfg["target_accuracy"],
+            estimated_quality=quality,
+            dataset_shape=f"{len(wc)} rows × {len(wc.columns)-1} columns",
+            rows_preserved_pct=rows_pct,
+            budget_remaining=max(budget, 0),
+            step_number=step,
+            max_steps=cfg["budget"],
+            active_session=self._active_session,
+            validate_calls_remaining=validate_calls_remaining(self._exploit),
+            done=False,
+            reward=round(reward, 4),
+        )
+
+    def _error_obs(self, msg: str) -> DataCentricObservation:
+        return DataCentricObservation(response=msg, done=False, reward=0.0)
+
+    def _unknown_cmd_obs(self, msg: str, step: int, budget: int,
+                         reward: float) -> DataCentricObservation:
+        return self._make_obs(
+            f"Unknown command: '{msg}'\n\n{AVAILABLE_COMMANDS}", step, budget, reward
+        )
+
+    @property
+    def state(self) -> State:
+        return self._state

server/dataset_generator.py

@@ -0,0 +1,257 @@
+"""
+Dataset Generator for Data-Centric RL Environment.
+
+Generates corrupted sklearn classification datasets with known ground truth.
+Each task has deterministic corruptions via seeded random.Random.
+
+CRITICAL: Always produces TWO copies:
+  ground_truth  → frozen, only read by grader
+  working_copy  → the only thing the agent can mutate
+"""
+
+import random
+from copy import deepcopy
+from typing import Any, Dict, Tuple
+
+import numpy as np
+import pandas as pd
+from sklearn.datasets import make_classification
+
+
+# ── Column metadata schema ──────────────────────────────────────────────────
+
+def _make_col_meta(expected_dtype: str, valid_range=None,
+                   valid_categories=None, is_nullable: bool = False) -> Dict:
+    return {
+        "expected_dtype": expected_dtype,
+        "valid_range": valid_range,
+        "valid_categories": valid_categories,
+        "is_nullable": is_nullable,
+    }
+
+
+# ── Task configurations ─────────────────────────────────────────────────────
+
+TASK_CONFIGS = {
+    "task_0_tutorial": {
+        "n_samples": 100,
+        "n_features": 4,
+        "n_classes": 2,
+        "n_informative": 3,
+        "budget": 30,
+        "target_accuracy": 0.73,
+        "baseline_accuracy": 0.62,
+        "description": "Single-issue tutorial. Fix missing values in 'age' to win.",
+    },
+    "task_1_easy": {
+        "n_samples": 200,
+        "n_features": 5,
+        "n_classes": 2,
+        "n_informative": 4,
+        "budget": 25,
+        "target_accuracy": 0.79,
+        "baseline_accuracy": 0.63,
+        "description": "Missing values + mild class imbalance.",
+    },
+    "task_2_medium": {
+        "n_samples": 500,
+        "n_features": 7,
+        "n_classes": 3,
+        "n_informative": 5,
+        "budget": 40,
+        "target_accuracy": 0.74,
+        "baseline_accuracy": 0.58,
+        "description": "Missing values, duplicates, class imbalance, type error.",
+    },
+    "task_3_hard": {
+        "n_samples": 900,
+        "n_features": 10,
+        "n_classes": 4,
+        "n_informative": 7,
+        "budget": 60,
+        "target_accuracy": 0.71,
+        "baseline_accuracy": 0.54,
+        "description": "Missing values, duplicates, imbalance, type errors, outliers, cross-column errors.",
+    },
+}
+
+
+# ── Generic feature names ───────────────────────────────────────────────────
+
+FEATURE_NAMES = ["age", "income", "score", "tenure", "balance",
+                 "transactions", "risk_level", "credit", "spend", "savings"]
+
+
+def _build_column_meta(feature_cols: list, task: str) -> Dict[str, Dict]:
+    meta = {}
+    for col in feature_cols:
+        meta[col] = _make_col_meta("float64", valid_range=(-10.0, 10.0))
+    # age gets tighter range for tutorial plausibility
+    if "age" in meta:
+        meta["age"] = _make_col_meta("float64", valid_range=(0.0, 100.0))
+    meta["target"] = _make_col_meta("int64", valid_categories=None)
+    return meta
+
+
+# ── Core generator ──────────────────────────────────────────────────────────
+
+def generate_dataset(task: str, seed: int = 42) -> Tuple[pd.DataFrame, pd.DataFrame, Dict[str, Any]]:
+    """
+    Generate a corrupted dataset for the given task.
+
+    Returns:
+        ground_truth  – clean DataFrame (frozen)
+        working_copy  – corrupted DataFrame (agent mutates this)
+        metadata      – task config + column metadata + original_length
+    """
+    cfg = TASK_CONFIGS[task]
+    rng = random.Random(seed)
+    np_rng = np.random.RandomState(seed)
+
+    n = cfg["n_samples"]
+    n_feat = cfg["n_features"]
+    n_cls = cfg["n_classes"]
+
+    # ── Generate clean classification data ──────────────────────────────────
+    X, y = make_classification(
+        n_samples=n,
+        n_features=n_feat,
+        n_informative=cfg["n_informative"],
+        n_redundant=max(0, n_feat - cfg["n_informative"] - 1),
+        n_classes=n_cls,
+        n_clusters_per_class=1,
+        weights=None,
+        random_state=seed,
+    )
+
+    cols = FEATURE_NAMES[:n_feat]
+    df_clean = pd.DataFrame(X, columns=cols)
+    df_clean["target"] = y
+
+    # Rescale 'age' column to [18, 80] for plausibility
+    if "age" in df_clean.columns:
+        mn, mx = df_clean["age"].min(), df_clean["age"].max()
+        df_clean["age"] = ((df_clean["age"] - mn) / (mx - mn + 1e-9)) * 62 + 18
+
+    ground_truth = deepcopy(df_clean)
+    working_copy = deepcopy(df_clean)
+
+    # ── Inject corruptions into working_copy only ────────────────────���───────
+    _inject_corruptions(working_copy, task, cfg, rng, np_rng, seed)
+
+    col_meta = _build_column_meta(cols, task)
+    metadata = {
+        **cfg,
+        "task": task,
+        "seed": seed,
+        "feature_cols": cols,
+        "col_meta": col_meta,
+        "original_length": len(working_copy),
+        "class_names": [str(c) for c in sorted(working_copy["target"].unique())],
+    }
+
+    return ground_truth, working_copy, metadata
+
+
+def _inject_corruptions(df: pd.DataFrame, task: str, cfg: dict,
+                        rng: random.Random, np_rng: np.random.RandomState,
+                        seed: int):
+    """Inject task-specific corruptions into df in-place."""
+
+    if task == "task_0_tutorial":
+        # Single issue: 20% missing in age only
+        _inject_missing(df, ["age"], frac=0.20, rng=rng)
+
+    elif task == "task_1_easy":
+        # Missing values 15% + mild class imbalance
+        cols = df.columns[:-1].tolist()
+        _inject_missing(df, cols[:2], frac=0.15, rng=rng)
+        _inject_class_imbalance(df, ratio=0.60, rng=rng, seed=seed)
+
+    elif task == "task_2_medium":
+        cols = df.columns[:-1].tolist()
+        _inject_missing(df, cols[:3], frac=0.12, rng=rng)
+        _inject_duplicates(df, frac=0.05, rng=rng)
+        _inject_class_imbalance(df, ratio=0.55, rng=rng, seed=seed)
+        _inject_type_error(df, cols[0], rng=rng, frac=0.04)
+
+    elif task == "task_3_hard":
+        cols = df.columns[:-1].tolist()
+        _inject_missing(df, cols[:4], frac=0.10, rng=rng)
+        _inject_duplicates(df, frac=0.05, rng=rng)
+        _inject_class_imbalance(df, ratio=0.50, rng=rng, seed=seed)
+        _inject_type_error(df, cols[0], rng=rng, frac=0.03)
+        _inject_outliers(df, cols[1], rng=rng, frac=0.03)
+        _inject_cross_column_errors(df, cols[2], cols[3], rng=rng, frac=0.02)
+
+
+def _inject_missing(df: pd.DataFrame, cols: list, frac: float, rng: random.Random):
+    for col in cols:
+        if col not in df.columns:
+            continue
+        indices = rng.sample(range(len(df)), int(len(df) * frac))
+        df.loc[indices, col] = np.nan
+
+
+def _inject_duplicates(df: pd.DataFrame, frac: float, rng: random.Random):
+    n_dups = max(1, int(len(df) * frac))
+    dup_indices = rng.choices(range(len(df)), k=n_dups)
+    dups = df.iloc[dup_indices].copy()
+    new_df = pd.concat([df, dups], ignore_index=True)
+    # Mutate the caller's DataFrame in-place by clearing and re-populating
+    df.drop(df.index, inplace=True)
+    df.drop(df.columns, axis=1, inplace=True)
+    for col in new_df.columns:
+        df[col] = new_df[col].values
+    df.reset_index(drop=True, inplace=True)
+
+
+def _inject_class_imbalance(df: pd.DataFrame, ratio: float,
+                             rng: random.Random, seed: int):
+    """Make class 0 account for `ratio` of rows, drop minority excess."""
+    target_col = "target"
+    classes = df[target_col].unique()
+    if len(classes) < 2:
+        return
+    major = int(classes[0])
+    n_major = int(len(df) * ratio)
+    major_idx = df[df[target_col] == major].index.tolist()
+    if len(major_idx) > n_major:
+        drop_n = len(major_idx) - n_major
+        to_drop = rng.sample(major_idx, drop_n)
+        df.drop(to_drop, inplace=True)
+        df.reset_index(drop=True, inplace=True)
+
+
+def _inject_type_error(df: pd.DataFrame, col: str, rng: random.Random, frac: float):
+    """Replace some float values with string 'ERR' to simulate type errors."""
+    if col not in df.columns:
+        return
+    indices = rng.sample(range(len(df)), max(1, int(len(df) * frac)))
+    df[col] = df[col].astype(object)
+    for i in indices:
+        df.at[i, col] = "ERR"
+
+
+def _inject_outliers(df: pd.DataFrame, col: str, rng: random.Random, frac: float):
+    if col not in df.columns:
+        return
+    indices = rng.sample(range(len(df)), max(1, int(len(df) * frac)))
+    for i in indices:
+        df.at[i, col] = rng.choice([999.0, -999.0])
+
+
+def _inject_cross_column_errors(df: pd.DataFrame, col_a: str, col_b: str,
+                                 rng: random.Random, frac: float):
+    """Make col_a < col_b for some rows (e.g. min > max violations)."""
+    if col_a not in df.columns or col_b not in df.columns:
+        return
+    indices = rng.sample(range(len(df)), max(1, int(len(df) * frac)))
+    for i in indices:
+        try:
+            a = float(df.at[i, col_a])
+            b = float(df.at[i, col_b])
+            if a >= b:
+                df.at[i, col_a], df.at[i, col_b] = b - 1.0, a + 1.0
+        except (ValueError, TypeError):
+            pass

server/grader.py

@@ -0,0 +1,400 @@
+"""
+Grader for Data-Centric RL Environment — using OpenEnv Rubric system.
+
+Implements 4 composable Rubric subclasses (nn.Module style, auto-registered
+as child rubrics) plus a root DataCentricRubric that aggregates them.
+
+Rubric hierarchy:
+    DataCentricRubric
+    ├── accuracy      : AccuracyRubric
+    ├── process       : ProcessRubric
+    ├── preservation  : PreservationRubric
+    └── efficiency    : EfficiencyRubric
+
+Also provides StepRubric for dense per-apply proxy feedback (no classifier).
+
+Backward-compatible: compute_*() free functions still work for existing callers.
+"""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import pandas as pd
+
+from openenv.core.rubrics.base import Rubric
+
+logger = logging.getLogger(__name__)
+
+# Must match openenv.yaml reward_range — enforced by DataCentricRubric
+REWARD_MIN: float = -1.0
+REWARD_MAX: float = 1.0
+
+
+# ── Lightweight quality score (no sklearn) ────────────────────────────────────
+
+def compute_lightweight_score(
+    working_copy: pd.DataFrame,
+    ground_truth: pd.DataFrame,
+    original_length: int,
+    col_meta: Dict,
+    initial_missing: int = None,
+) -> float:
+    """
+    Fast quality score comparing working_copy to ground_truth structure.
+    Does NOT run sklearn — used for dense per-step feedback.
+
+    Score is [0.0, 1.0] composed of:
+      - missing value reduction (40%)
+      - duplicate reduction     (20%)
+      - type correctness        (20%)
+      - row preservation        (20%)
+    """
+    score = 0.0
+
+    # 1. Missing value reduction
+    wc_missing = int(working_copy.isnull().sum().sum())
+    denom = initial_missing if (initial_missing is not None and initial_missing > 0) else max(wc_missing, 1)
+    missing_score = max(0.0, 1.0 - wc_missing / denom) if denom > 0 else 1.0
+    score += 0.40 * missing_score
+
+    # 2. Duplicate reduction
+    n_dups_wc = int(working_copy.duplicated().sum())
+    n_dups_gt = int(ground_truth.duplicated().sum())
+    if n_dups_gt == 0 and n_dups_wc == 0:
+        dup_score = 1.0
+    elif n_dups_gt == 0:
+        dup_score = max(0.0, 1.0 - n_dups_wc / max(len(working_copy), 1))
+    else:
+        dup_score = max(0.0, 1.0 - n_dups_wc / max(n_dups_gt, 1))
+    score += 0.20 * dup_score
+
+    # 3. Type correctness
+    type_ok, type_total = 0, 0
+    for col, meta in col_meta.items():
+        if col == "target" or col not in working_copy.columns:
+            continue
+        if meta.get("expected_dtype", "float64") in ("float64", "int64"):
+            type_total += 1
+            err_count = sum(
+                1 for val in working_copy[col].dropna()
+                if not _can_float(val)
+            )
+            if err_count == 0:
+                type_ok += 1
+    score += 0.20 * ((type_ok / type_total) if type_total > 0 else 1.0)
+
+    # 4. Row preservation
+    score += 0.20 * min(len(working_copy) / max(original_length, 1), 1.0)
+
+    return round(min(score, 1.0), 4)
+
+
+def _can_float(val: Any) -> bool:
+    try:
+        float(val)
+        return True
+    except (ValueError, TypeError):
+        return False
+
+
+# ── Rubric 1: Accuracy ────────────────────────────────────────────────────────
+
+class AccuracyRubric(Rubric):
+    """
+    Main RL signal: rewards accuracy improvement, penalises regression.
+    Adds a large terminal bonus when the agent submits and crosses the target.
+    """
+
+    def forward(self, action: Any, observation: Any) -> float:
+        current   = observation.get("current_accuracy", 0.0)
+        previous  = observation.get("previous_accuracy", 0.0)
+        baseline  = observation.get("baseline_accuracy", 0.0)
+        target    = observation.get("target_accuracy", 0.80)
+        is_submit = str(action).strip().lower() == "submit"
+
+        improvement = current - previous
+        if improvement > 0:
+            reward = improvement * 2.5
+        elif improvement < 0:
+            reward = improvement * 2.0   # regression penalised harder
+        else:
+            reward = 0.0
+
+        if is_submit:
+            if current >= target:
+                reward += 0.50            # big terminal bonus
+            else:
+                progress_range = target - baseline
+                if progress_range > 0:
+                    progress = (current - baseline) / progress_range
+                    reward += max(0.0, progress) * 0.25
+
+        logger.debug("AccuracyRubric: imp=%.4f reward=%.4f submit=%s", improvement, reward, is_submit)
+        return round(reward, 4)
+
+
+# ── Rubric 2: Process ─────────────────────────────────────────────────────────
+
+class ProcessRubric(Rubric):
+    """
+    Rewards smart workflow patterns (inspect → query → apply → validate).
+    Penalises blind apply-without-query and submit-without-validate.
+    """
+
+    def forward(self, action: Any, observation: Any) -> float:
+        history: List[str] = observation.get("action_history", [])
+        current_action = str(action)
+        full_history = (history + [current_action])[-5:]
+        reward = 0.0
+
+        def _cmd(a: str) -> str:
+            return a.split()[0].lower()
+
+        cmd = _cmd(current_action)
+        prev_cmds = [_cmd(h) for h in full_history[:-1][-3:]]
+
+        if cmd.startswith("query_"):
+            if "inspect_dataset" in prev_cmds or "inspect_model" in prev_cmds:
+                reward += 0.02
+
+        if cmd == "apply":
+            if any(p.startswith("query_") for p in prev_cmds):
+                reward += 0.05
+            else:
+                reward -= 0.04
+
+        if cmd == "validate" and "apply" in prev_cmds:
+            reward += 0.03
+
+        if cmd == "reject":
+            reward += 0.01
+
+        if cmd == "submit":
+            all_cmds = [_cmd(h) for h in history]
+            if "validate" not in all_cmds:
+                reward -= 0.10
+
+        logger.debug("ProcessRubric: action=%s reward=%.4f", current_action, reward)
+        return round(reward, 4)
+
+
+# ── Rubric 3: Preservation ────────────────────────────────────────────────────
+
+class PreservationRubric(Rubric):
+    """
+    Rewards row preservation. Independent of accuracy — prevents the agent
+    from 'cheating' by deleting rows to inflate classifier confidence.
+    """
+
+    def forward(self, action: Any, observation: Any) -> float:
+        current_rows  = observation.get("current_rows", 0)
+        original_rows = observation.get("original_rows", 1)
+        rows_preserved = current_rows / max(original_rows, 1)
+
+        if rows_preserved >= 0.90:
+            reward = 0.05
+        elif rows_preserved >= 0.80:
+            reward = 0.02
+        elif rows_preserved >= 0.70:
+            reward = 0.00
+        elif rows_preserved >= 0.50:
+            reward = -0.10
+        else:
+            reward = -0.40
+
+        logger.debug("PreservationRubric: pct=%.2f reward=%.4f", rows_preserved, reward)
+        return round(reward, 4)
+
+
+# ── Rubric 4: Efficiency ──────────────────────────────────────────────────────
+
+class EfficiencyRubric(Rubric):
+    """
+    Computed ONLY at submit. Rewards high accuracy gain per budget step used.
+    Encourages the agent to be surgical rather than spray-and-pray.
+    """
+
+    def forward(self, action: Any, observation: Any) -> float:
+        if str(action).strip().lower() != "submit":
+            return 0.0
+
+        baseline         = observation.get("baseline_accuracy", 0.0)
+        current          = observation.get("current_accuracy", 0.0)
+        original_budget  = observation.get("original_budget", 1)
+        budget_remaining = observation.get("budget_remaining", 0)
+        budget_used      = max(original_budget - budget_remaining, 1)
+        accuracy_gain    = current - baseline
+
+        if accuracy_gain <= 0:
+            reward = -0.05
+        else:
+            reward = min((accuracy_gain / budget_used) * 3.0, 0.20)
+
+        logger.debug("EfficiencyRubric: gain=%.4f used=%d reward=%.4f",
+                     accuracy_gain, budget_used, reward)
+        return round(reward, 4)
+
+
+# ── Rubric 5: Step (proxy, no classifier) ────────────────────────────────────
+
+class StepRubric(Rubric):
+    """
+    Dense per-apply proxy reward — does NOT run the RF classifier.
+    Uses lightweight quality score delta to give feedback between validate calls.
+    Registered as a standalone rubric (not a child of DataCentricRubric)
+    because it fires on every apply step, not just at episode end.
+    """
+
+    def forward(self, action: Any, observation: Any) -> float:
+        if not str(action).startswith("apply"):
+            return 0.0
+
+        q_before = observation.get("quality_before", 0.0)
+        q_after  = observation.get("quality_after", 0.0)
+        rows_pct = observation.get("rows_preserved_after", 1.0)
+
+        r = float(np.clip((q_after - q_before) * 0.3, -0.20, 0.10))
+
+        if rows_pct >= 0.95:
+            r += 0.02
+        elif rows_pct >= 0.90:
+            r += 0.01
+        elif rows_pct < 0.80:
+            r -= 0.10
+
+        return float(np.clip(r, -0.30, 0.15))
+
+
+# ── Root Rubric (aggregates all components) ───────────────────────────────────
+
+class DataCentricRubric(Rubric):
+    """
+    Root composable rubric for the Data-Centric AI environment.
+
+    Child rubrics are auto-registered (PyTorch nn.Module style):
+        rubric.accuracy     → AccuracyRubric
+        rubric.process      → ProcessRubric
+        rubric.preservation → PreservationRubric
+        rubric.efficiency   → EfficiencyRubric
+
+    Call rubric(action, obs_dict) to get total clamped reward [-1.0, 1.0].
+    Inspect rubric.accuracy.last_score etc. for per-component breakdown.
+    """
+
+    def __init__(self):
+        super().__init__()
+        # Assigned as attributes — auto-registered as children by Rubric.__setattr__
+        self.accuracy    = AccuracyRubric()
+        self.process     = ProcessRubric()
+        self.preservation = PreservationRubric()
+        self.efficiency  = EfficiencyRubric()
+
+    def forward(self, action: Any, observation: Any) -> float:
+        r_acc  = self.accuracy(action, observation)
+        r_proc = self.process(action, observation)
+        r_pres = self.preservation(action, observation)
+        r_eff  = self.efficiency(action, observation)
+        total  = r_acc + r_proc + r_pres + r_eff
+
+        clamped = float(np.clip(total, REWARD_MIN, REWARD_MAX))
+        if __debug__ and abs(clamped - total) > 1e-6:
+            logger.warning("Reward %.4f clamped → %.4f", total, clamped)
+
+        logger.info(
+            "REWARD | accuracy=%.4f process=%.4f preservation=%.4f "
+            "efficiency=%.4f TOTAL=%.4f (clamped=%.4f)",
+            r_acc, r_proc, r_pres, r_eff, total, clamped,
+        )
+        return round(clamped, 4)
+
+    def breakdown(self) -> Dict[str, Optional[float]]:
+        """Return last_score for each child rubric — useful for logging."""
+        return {
+            "accuracy":     self.accuracy.last_score,
+            "process":      self.process.last_score,
+            "preservation": self.preservation.last_score,
+            "efficiency":   self.efficiency.last_score,
+        }
+
+
+# ── Singleton — reuse across episode steps ────────────────────────────────────
+
+_rubric: Optional[DataCentricRubric] = None
+_step_rubric: Optional[StepRubric] = None
+
+
+def get_rubric() -> DataCentricRubric:
+    global _rubric
+    if _rubric is None:
+        _rubric = DataCentricRubric()
+    return _rubric
+
+
+def get_step_rubric() -> StepRubric:
+    global _step_rubric
+    if _step_rubric is None:
+        _step_rubric = StepRubric()
+    return _step_rubric
+
+
+# ── Backward-compatible free functions ───────────────────────────────────────
+# (called by data_centric_environment.py — no changes needed there)
+
+def compute_accuracy_reward(
+    current_accuracy: float, previous_accuracy: float,
+    baseline_accuracy: float, target_accuracy: float,
+    is_submit: bool = False,
+) -> float:
+    obs = dict(current_accuracy=current_accuracy, previous_accuracy=previous_accuracy,
+               baseline_accuracy=baseline_accuracy, target_accuracy=target_accuracy)
+    action = "submit" if is_submit else "step"
+    return get_rubric().accuracy(action, obs)
+
+
+def compute_process_reward(action_history: List[str], current_action: str) -> float:
+    obs = dict(action_history=action_history)
+    return get_rubric().process(current_action, obs)
+
+
+def compute_preservation_reward(current_rows: int, original_rows: int) -> float:
+    obs = dict(current_rows=current_rows, original_rows=original_rows)
+    return get_rubric().preservation("step", obs)
+
+
+def compute_efficiency_reward(
+    current_accuracy: float, baseline_accuracy: float,
+    original_budget: int, budget_remaining: int,
+) -> float:
+    obs = dict(current_accuracy=current_accuracy, baseline_accuracy=baseline_accuracy,
+               original_budget=original_budget, budget_remaining=budget_remaining)
+    return get_rubric().efficiency("submit", obs)
+
+
+def compute_step_reward(
+    action: str, quality_before: float, quality_after: float,
+    rows_preserved_after: float,
+) -> float:
+    obs = dict(quality_before=quality_before, quality_after=quality_after,
+               rows_preserved_after=rows_preserved_after)
+    return get_step_rubric()(action, obs)
+
+
+def compute_total_reward(
+    reward_accuracy: float,
+    reward_process: float,
+    reward_preservation: float,
+    reward_efficiency: float = 0.0,
+    reward_step: float = 0.0,
+) -> float:
+    total = reward_accuracy + reward_process + reward_preservation + reward_efficiency + reward_step
+    clamped = float(np.clip(total, REWARD_MIN, REWARD_MAX))
+    if __debug__ and abs(clamped - total) > 1e-6:
+        logger.warning("Reward %.4f clamped → %.4f", total, clamped)
+    logger.info(
+        "REWARD BREAKDOWN: accuracy=%.4f process=%.4f preservation=%.4f "
+        "efficiency=%.4f step=%.4f TOTAL=%.4f (clamped=%.4f)",
+        reward_accuracy, reward_process, reward_preservation,
+        reward_efficiency, reward_step, total, clamped,
+    )
+    return round(clamped, 4)

server/model_evaluator.py

@@ -0,0 +1,228 @@
+"""
+Model Evaluator for Data-Centric RL Environment.
+
+Uses sklearn RandomForestClassifier (primary) + LogisticRegression (secondary)
+with hash-based caching. Trains on working_copy, evaluates on held-out
+ground_truth test split.
+
+Primary accuracy (RF) drives all rewards and grading.
+Secondary accuracy (LR) is diagnostic — shows whether improvements generalise
+beyond the RF decision boundary (overfitting detection).
+"""
+
+import hashlib
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+
+
+class ModelEvaluator:
+    """Caching dual-classifier evaluator.
+
+    Args:
+        fast_mode: If True, uses n_estimators=20 (for GRPO rollouts, ~4x faster).
+                   If False, uses n_estimators=100 (for final eval, more accurate).
+    """
+
+    def __init__(self, fast_mode: bool = False):
+        self._cache_hash: Optional[str] = None
+        self._cached_accuracy: float = 0.0
+        self._cached_per_class: Dict = {}
+        self._cached_lr_accuracy: float = 0.0
+        self._cached_feature_importance: Dict[str, float] = {}
+        self._cached = False
+        self._fast_mode = fast_mode
+        self._n_estimators = 20 if fast_mode else 100
+
+    def _compute_hash(self, df: pd.DataFrame) -> str:
+        try:
+            return hashlib.md5(
+                pd.util.hash_pandas_object(df, index=True).values.tobytes()
+            ).hexdigest()
+        except Exception:
+            return hashlib.md5(df.to_json().encode()).hexdigest()
+
+    def evaluate(
+        self,
+        working_copy: pd.DataFrame,
+        ground_truth: pd.DataFrame,
+        test_size: float = 0.25,
+        seed: int = 42,
+    ) -> Tuple[float, Dict, bool, float]:
+        """
+        Train on working_copy; evaluate on held-out ground_truth test split.
+
+        Returns:
+            rf_accuracy    – float  (primary, used for rewards)
+            per_class      – dict from classification_report
+            from_cache     – True if result came from cache (no retrain)
+            lr_accuracy    – float  (secondary, diagnostic only)
+        """
+        from sklearn.ensemble import RandomForestClassifier
+        from sklearn.linear_model import LogisticRegression
+        from sklearn.metrics import classification_report
+        from sklearn.model_selection import train_test_split
+
+        current_hash = self._compute_hash(working_copy)
+
+        if self._cached and current_hash == self._cache_hash:
+            return self._cached_accuracy, self._cached_per_class, True, self._cached_lr_accuracy
+
+        # Prepare train set from working_copy
+        wc = working_copy.copy()
+        wc = wc.dropna(subset=["target"])
+        for col in wc.columns:
+            if col != "target":
+                wc[col] = pd.to_numeric(wc[col], errors="coerce")
+        wc = wc.dropna()
+
+        _empty = (0.0, {}, False, 0.0)
+        if len(wc) < 10:
+            self._cache_hash = current_hash
+            self._cached_accuracy = 0.0
+            self._cached_per_class = {}
+            self._cached_lr_accuracy = 0.0
+            self._cached_feature_importance = {}
+            self._cached = True
+            return _empty
+
+        X_train = wc.drop("target", axis=1).values
+        y_train = wc["target"].astype(int).values
+
+        # Build test set from ground_truth
+        gt = ground_truth.copy().dropna()
+        for col in gt.columns:
+            if col != "target":
+                gt[col] = pd.to_numeric(gt[col], errors="coerce")
+        gt = gt.dropna()
+
+        if len(gt) < 10:
+            self._cache_hash = current_hash
+            self._cached_accuracy = 0.0
+            self._cached_per_class = {}
+            self._cached_lr_accuracy = 0.0
+            self._cached_feature_importance = {}
+            self._cached = True
+            return _empty
+
+        _, X_test_df, _, y_test_df = train_test_split(
+            gt.drop("target", axis=1),
+            gt["target"],
+            test_size=test_size,
+            random_state=seed,
+            stratify=gt["target"] if gt["target"].nunique() > 1 else None,
+        )
+        y_test = y_test_df.astype(int).values
+
+        # Align columns
+        train_cols = list(wc.drop("target", axis=1).columns)
+        test_cols = list(X_test_df.columns)
+        shared = [c for c in train_cols if c in test_cols]
+        if not shared:
+            self._cache_hash = current_hash
+            self._cached_accuracy = 0.0
+            self._cached_per_class = {}
+            self._cached_lr_accuracy = 0.0
+            self._cached_feature_importance = {}
+            self._cached = True
+            return _empty
+
+        X_train_arr = wc[shared].values
+        X_test_arr = X_test_df[shared].values
+
+        # ── Primary: Random Forest ──────────────────────────────────────────
+        rf = RandomForestClassifier(
+            n_estimators=self._n_estimators,
+            random_state=42,
+            n_jobs=1,
+        )
+        rf.fit(X_train_arr, y_train)
+        y_pred_rf = rf.predict(X_test_arr)
+
+        rf_accuracy = float(rf.score(X_test_arr, y_test))
+        try:
+            per_class = classification_report(
+                y_test, y_pred_rf, output_dict=True, zero_division=0
+            )
+        except Exception:
+            per_class = {}
+
+        # Feature importance (RF gives this for free)
+        feature_importance: Dict[str, float] = {}
+        if hasattr(rf, "feature_importances_"):
+            importances = rf.feature_importances_
+            for col, imp in zip(shared, importances):
+                feature_importance[col] = round(float(imp), 4)
+
+        # ── Secondary: Logistic Regression (diagnostic) ─────────────────────
+        lr_accuracy = 0.0
+        if not self._fast_mode:  # skip LR in fast_mode to keep GRPO rollouts quick
+            try:
+                lr = LogisticRegression(max_iter=500, random_state=42, n_jobs=1)
+                lr.fit(X_train_arr, y_train)
+                lr_accuracy = float(lr.score(X_test_arr, y_test))
+            except Exception:
+                lr_accuracy = 0.0
+        else:
+            # In fast_mode, reuse RF accuracy as placeholder (not shown to agent)
+            lr_accuracy = rf_accuracy
+
+        # Update cache
+        self._cache_hash = current_hash
+        self._cached_accuracy = rf_accuracy
+        self._cached_per_class = per_class
+        self._cached_lr_accuracy = lr_accuracy
+        self._cached_feature_importance = feature_importance
+        self._cached = True
+
+        return rf_accuracy, per_class, False, lr_accuracy
+
+    def agreement_signal(self, rf_acc: float, lr_acc: float,
+                         prev_rf: float, prev_lr: float) -> str:
+        """
+        Compare RF vs LR improvement direction.
+        Returns a signal string for the agent to reason about.
+        """
+        rf_improved = rf_acc > prev_rf + 0.005
+        lr_improved = lr_acc > prev_lr + 0.005
+        rf_declined = rf_acc < prev_rf - 0.005
+        lr_declined = lr_acc < prev_lr - 0.005
+
+        if rf_improved and lr_improved:
+            return "BOTH_AGREE_IMPROVE — fix is robust and generalises"
+        elif rf_improved and lr_declined:
+            return "DISAGREE — RF improved but LR declined (possible RF-specific overfitting)"
+        elif rf_declined and lr_declined:
+            return "BOTH_DECLINED — last change hurt both classifiers, consider undo"
+        elif not rf_improved and not rf_declined:
+            return "NO_CHANGE — last operation had no measurable effect"
+        else:
+            return "MIXED — marginal changes, continue and validate again"
+
+    def feature_importance_text(self, top_n: int = 5) -> str:
+        """Return formatted feature importance string for agent observation."""
+        if not self._cached_feature_importance:
+            return ""
+        sorted_feats = sorted(
+            self._cached_feature_importance.items(),
+            key=lambda x: -x[1]
+        )[:top_n]
+        parts = [f"{col} ({imp:.3f})" for col, imp in sorted_feats]
+        return "Feature importance: " + " > ".join(parts)
+
+    def invalidate_cache(self):
+        self._cached = False
+        self._cache_hash = None
+
+    @property
+    def last_accuracy(self) -> float:
+        return self._cached_accuracy if self._cached else 0.0
+
+    @property
+    def last_lr_accuracy(self) -> float:
+        return self._cached_lr_accuracy if self._cached else 0.0
+
+    @property
+    def last_feature_importance(self) -> Dict[str, float]:
+        return self._cached_feature_importance if self._cached else {}

server/requirements.txt

@@ -0,0 +1,6 @@
+openenv-core>=0.2.0
+fastapi>=0.115.0
+uvicorn>=0.24.0
+scikit-learn>=1.3.0
+pandas>=2.0.0
+numpy>=1.24.0

server/specialist_agents.py

@@ -0,0 +1,654 @@
+"""
+Specialist Agents for Data-Centric RL Environment.
+
+Agents:
+  CleanerAgent    — detects missing values, duplicates, type errors
+  AugmenterAgent  — suggests synthetic minority class samples
+  BalancerAgent   — recommends resampling strategies
+  ValidatorAgent  — checks column metadata rule violations (costs 2 budget)
+  AnalystAgent    — holistic diagnosis + prioritised action plan (costs 1 budget)
+
+Also exports:
+  compute_drift() — per-column distribution drift score (no scipy needed)
+"""
+
+import hashlib
+import random
+import uuid
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import pandas as pd
+
+
+# ── Recommendation / Violation dataclasses ──────────────────────────────────
+
+@dataclass
+class Recommendation:
+    id: int
+    description: str
+    action_type: str
+    estimated_impact: float
+    confidence: float
+    session_id: str
+    _payload: Dict[str, Any] = field(default_factory=dict, repr=False)
+
+
+@dataclass
+class Violation:
+    column: str
+    rule: str
+    count: int
+    description: str
+    severity: str = "WARNING"   # NEW: CRITICAL / WARNING / INFO
+
+
+# ── Session registry ─────────────────────────────────────────────────────────
+
+class SessionRegistry:
+    """Tracks the active recommendation session to detect stale IDs."""
+
+    def __init__(self):
+        self.current_session_id: str = ""
+        self.recommendations: Dict[int, Recommendation] = {}
+
+    def new_session(self) -> str:
+        self.current_session_id = str(uuid.uuid4())
+        self.recommendations = {}
+        return self.current_session_id
+
+    def register(self, recs: List[Recommendation]) -> None:
+        for r in recs:
+            self.recommendations[r.id] = r
+
+    def get(self, rec_id: int, session_id: str) -> Optional[Recommendation]:
+        if session_id != self.current_session_id:
+            return None
+        return self.recommendations.get(rec_id)
+
+    def is_valid_session(self, session_id: str) -> bool:
+        return session_id == self.current_session_id
+
+
+# ── Shared helpers ───────────────────────────────────────────────────────────
+
+def _seeded_rng(df: pd.DataFrame, salt: str = "") -> random.Random:
+    h = hashlib.md5((df.to_json() + salt).encode()).hexdigest()
+    return random.Random(int(h[:8], 16))
+
+
+def _col_stats(series: pd.Series) -> Dict[str, float]:
+    """Return basic stats dict for a numeric series."""
+    s = pd.to_numeric(series, errors="coerce").dropna()
+    if len(s) == 0:
+        return {"mean": 0.0, "median": 0.0, "std": 0.0, "skew": 0.0, "n": 0}
+    return {
+        "mean": float(s.mean()),
+        "median": float(s.median()),
+        "std": float(s.std()) if len(s) > 1 else 0.0,
+        "skew": float(s.skew()) if len(s) > 2 else 0.0,
+        "n": len(s),
+    }
+
+
+def _impute_strategy(stats: Dict[str, float]) -> tuple:
+    """Choose mean vs median based on skewness. Returns (strategy, value, reason)."""
+    skew = abs(stats["skew"])
+    if skew > 1.0:
+        return "median", stats["median"], f"right-skewed (skew={stats['skew']:.2f}), median more robust"
+    elif skew > 0.5:
+        return "median", stats["median"], f"moderately skewed (skew={stats['skew']:.2f}), median preferred"
+    else:
+        return "mean", stats["mean"], f"near-symmetric (skew={stats['skew']:.2f}), mean appropriate"
+
+
+# ── Drift Detection ──────────────────────────────────────────────────────────
+
+def compute_drift(working_copy: pd.DataFrame, ground_truth: pd.DataFrame) -> Dict[str, float]:
+    """
+    Per-column distribution drift score comparing working_copy to ground_truth.
+    Uses mean-shift + std-ratio (no scipy dependency).
+    Returns dict: column -> drift_score (0.0 = no drift, >1.0 = HIGH drift).
+    """
+    drift = {}
+    for col in working_copy.columns:
+        if col == "target":
+            continue
+        try:
+            wc_vals = pd.to_numeric(working_copy[col], errors="coerce").dropna()
+            gt_vals = pd.to_numeric(ground_truth[col], errors="coerce").dropna()
+            if len(wc_vals) == 0 or len(gt_vals) == 0:
+                drift[col] = 0.0
+                continue
+            mean_shift = abs(wc_vals.mean() - gt_vals.mean()) / (gt_vals.std() + 1e-8)
+            std_ratio = wc_vals.std() / (gt_vals.std() + 1e-8)
+            drift[col] = round(float(mean_shift + abs(1.0 - std_ratio)), 3)
+        except Exception:
+            drift[col] = 0.0
+    return drift
+
+
+def _drift_label(score: float) -> str:
+    if score < 0.2:
+        return "NONE"
+    elif score < 0.5:
+        return "LOW"
+    elif score < 1.0:
+        return "MEDIUM"
+    else:
+        return "HIGH"
+
+
+def format_drift_summary(drift: Dict[str, float]) -> str:
+    """Return one-line drift summary for agent observation."""
+    if not drift:
+        return ""
+    parts = [f"{col} ({_drift_label(v)})" for col, v in sorted(drift.items(), key=lambda x: -x[1])]
+    return "Distribution drift: " + " | ".join(parts[:5])  # top 5 most drifted
+
+
+# ── CleanerAgent ─────────────────────────────────────────────────────────────
+
+class CleanerAgent:
+    """
+    Analyses working_copy for missing values, duplicates, type mismatches.
+    Returns 2-4 recommendations with statistical reasoning.
+
+    Hidden flaw (15% of calls, deterministic): occasionally recommends
+    removing rows that are valid. Detectable because estimated_impact < 0.
+    """
+
+    def query(self, df: pd.DataFrame, session_registry: SessionRegistry,
+              col_meta: Dict) -> List[Recommendation]:
+        sid = session_registry.new_session()
+        rng = _seeded_rng(df, "cleaner")
+        recs = []
+        rec_id = 1
+        n_rows = len(df)
+
+        # --- Missing value recommendations with statistical reasoning ---
+        for col in df.columns:
+            if col == "target":
+                continue
+            n_missing = int(df[col].isna().sum())
+            if n_missing == 0:
+                continue
+
+            pct_missing = n_missing / n_rows * 100
+            stats = _col_stats(df[col])
+            strategy, value, reason = _impute_strategy(stats)
+
+            # Confidence: lower if >30% missing (imputation less reliable)
+            confidence = round(max(0.60, 0.92 - (pct_missing / 100) * 0.5), 2)
+
+            # Risk label
+            if pct_missing < 5:
+                risk = "LOW"
+            elif pct_missing < 20:
+                risk = "MEDIUM"
+            else:
+                risk = "HIGH — imputation may introduce bias"
+
+            mean_median_delta = abs(stats["mean"] - stats["median"])
+            description = (
+                f"Fill {n_missing}/{n_rows} ({pct_missing:.1f}%) missing values in '{col}' "
+                f"using {strategy} ({value:.2f}). "
+                f"Reason: {reason}. "
+                f"Mean={stats['mean']:.2f}, Median={stats['median']:.2f} "
+                f"(delta={mean_median_delta:.2f}). Risk: {risk}."
+            )
+
+            recs.append(Recommendation(
+                id=rec_id,
+                description=description,
+                action_type="fill_missing",
+                estimated_impact=round(min(0.03 + pct_missing / 100 * 0.3, 0.12), 3),
+                confidence=confidence,
+                session_id=sid,
+                _payload={"action": "fill_missing", "column": col, "strategy": strategy},
+            ))
+            rec_id += 1
+
+        # --- Duplicate recommendation ---
+        n_dups = int(df.duplicated().sum())
+        if n_dups > 0:
+            pct_dups = n_dups / n_rows * 100
+            recs.append(Recommendation(
+                id=rec_id,
+                description=(
+                    f"Remove {n_dups} duplicate rows ({pct_dups:.1f}% of dataset). "
+                    f"Duplicates bias the classifier toward overrepresented patterns. Risk: LOW."
+                ),
+                action_type="remove_duplicates",
+                estimated_impact=round(min(0.02 + n_dups / n_rows * 0.15, 0.08), 3),
+                confidence=0.92,
+                session_id=sid,
+                _payload={"action": "remove_duplicates"},
+            ))
+            rec_id += 1
+
+        # --- Type error recommendations ---
+        for col in df.columns:
+            if col == "target":
+                continue
+            meta = col_meta.get(col, {})
+            expected = meta.get("expected_dtype", "float64")
+            if expected in ("float64", "int64"):
+                n_errors = sum(1 for val in df[col].dropna()
+                               if not _is_numeric(val))
+                if n_errors > 0:
+                    pct_err = n_errors / n_rows * 100
+                    recs.append(Recommendation(
+                        id=rec_id,
+                        description=(
+                            f"Fix {n_errors} type errors ({pct_err:.1f}%) in '{col}' "
+                            f"(non-numeric values coerced to NaN, then filled with mean). "
+                            f"Expected dtype: {expected}. Risk: LOW."
+                        ),
+                        action_type="fix_type_errors",
+                        estimated_impact=round(min(0.04 + n_errors / n_rows * 0.2, 0.10), 3),
+                        confidence=0.88,
+                        session_id=sid,
+                        _payload={"action": "fix_type_errors", "column": col},
+                    ))
+                    rec_id += 1
+
+        # --- Hidden flaw: ~15% chance of recommending valid row removal ---
+        flaw_hash = int(hashlib.md5(df.to_json().encode()).hexdigest()[:4], 16)
+        if flaw_hash % 100 < 15 and len(recs) < 4:
+            col = rng.choice([c for c in df.columns if c != "target"])
+            recs.append(Recommendation(
+                id=rec_id,
+                description=(
+                    f"Remove rows where '{col}' is below the 5th percentile "
+                    f"(suspected outliers). Confidence LOW — verify with query_validator first."
+                ),
+                action_type="remove_outlier_rows",
+                estimated_impact=round(rng.uniform(-0.05, 0.01), 3),
+                confidence=round(rng.uniform(0.55, 0.70), 2),
+                session_id=sid,
+                _payload={"action": "remove_outlier_rows", "column": col, "pct": 5},
+            ))
+
+        # Keep top 4 by estimated_impact, re-number
+        recs = sorted(recs, key=lambda r: -r.estimated_impact)[:4]
+        for i, r in enumerate(recs, 1):
+            r.id = i
+        session_registry.register(recs)
+        return recs
+
+
+def _is_numeric(val) -> bool:
+    try:
+        float(val)
+        return True
+    except (ValueError, TypeError):
+        return False
+
+
+# ── AugmenterAgent ───────────────────────────────────────────────────────────
+
+class AugmenterAgent:
+    """
+    Detects underrepresented classes and suggests synthetic samples.
+    Returns 1-3 recommendations with class distribution reasoning.
+
+    Hidden flaw: sometimes suggests out-of-distribution samples (flagged).
+    """
+
+    def query(self, df: pd.DataFrame, session_registry: SessionRegistry,
+              class_name: Optional[str] = None) -> List[Recommendation]:
+        sid = session_registry.new_session()
+        rng = _seeded_rng(df, "augmenter")
+        value_counts = df["target"].value_counts()
+        total = len(df)
+        recs = []
+        rec_id = 1
+
+        targets = [class_name] if class_name else [str(c) for c in value_counts.index]
+
+        # Class distribution context
+        dist_str = ", ".join(f"class {k}: {v} ({v/total*100:.1f}%)"
+                             for k, v in value_counts.items())
+
+        for cls in targets[:3]:
+            try:
+                cls_int = int(cls)
+            except (ValueError, TypeError):
+                continue
+            if cls_int not in value_counts.index:
+                continue
+
+            count = value_counts[cls_int]
+            max_count = value_counts.max()
+            gap = max_count - count
+            if gap <= 0:
+                continue
+
+            n_synth = min(gap, max(5, int(gap * 0.5)))
+            ratio_before = count / max_count
+            ratio_after = (count + n_synth) / max_count
+
+            flaw_hash = int(hashlib.md5((df.to_json() + cls).encode()).hexdigest()[:4], 16)
+            is_ood = flaw_hash % 100 < 20
+            impact = round(min(0.04 + n_synth / total * 0.4, 0.10), 3)
+            if is_ood:
+                impact = round(rng.uniform(-0.02, 0.02), 3)
+
+            ood_note = " [WARNING: high OOD risk — run query_validator before applying]" if is_ood else ""
+            risk = "HIGH" if is_ood else ("MEDIUM" if ratio_before < 0.3 else "LOW")
+
+            description = (
+                f"Generate {n_synth} synthetic samples for class '{cls}' via Gaussian perturbation. "
+                f"Distribution: {dist_str}. "
+                f"Imbalance ratio before: {ratio_before:.2f} → after: {ratio_after:.2f}. "
+                f"Risk: {risk}.{ood_note}"
+            )
+
+            recs.append(Recommendation(
+                id=rec_id,
+                description=description,
+                action_type="augment_class",
+                estimated_impact=impact,
+                confidence=round(0.60 if is_ood else 0.82, 2),
+                session_id=sid,
+                _payload={
+                    "action": "augment_class",
+                    "class": cls_int,
+                    "n_synth": n_synth,
+                    "ood": is_ood,
+                },
+            ))
+            rec_id += 1
+
+        session_registry.register(recs)
+        return recs
+
+
+# ── BalancerAgent ─────────────────────────────────────────────────────────────
+
+class BalancerAgent:
+    """
+    Recommends resampling strategies for class imbalance.
+    Returns 1-2 recommendations with entropy and ratio reasoning.
+
+    Hidden flaw: occasionally over-balances (minority becomes too large).
+    """
+
+    def query(self, df: pd.DataFrame, session_registry: SessionRegistry) -> List[Recommendation]:
+        sid = session_registry.new_session()
+        rng = _seeded_rng(df, "balancer")
+        value_counts = df["target"].value_counts()
+        recs = []
+        rec_id = 1
+
+        if len(value_counts) < 2:
+            session_registry.register([])
+            return []
+
+        min_cls = int(value_counts.idxmin())
+        max_cls = int(value_counts.idxmax())
+        min_count = int(value_counts.min())
+        max_count = int(value_counts.max())
+        imbalance_ratio = min_count / max_count
+
+        # Class distribution entropy (0=perfectly imbalanced, 1=perfectly balanced)
+        probs = value_counts / value_counts.sum()
+        entropy = float(-np.sum(probs * np.log2(probs + 1e-9)))
+        max_entropy = np.log2(len(value_counts))
+        entropy_pct = entropy / max_entropy * 100 if max_entropy > 0 else 0
+
+        flaw_hash = int(hashlib.md5(df.to_json().encode()).hexdigest()[:4], 16)
+        is_overbalance = flaw_hash % 100 < 20
+        target_count = max_count if not is_overbalance else int(max_count * 1.5)
+        ratio_after = min(1.0, min_count / target_count) if target_count > 0 else 1.0
+
+        overbalance_note = (
+            " [WARNING: target exceeds majority class size — may over-correct and hurt generalisation]"
+            if is_overbalance else ""
+        )
+        risk = "HIGH" if is_overbalance else ("MEDIUM" if imbalance_ratio < 0.3 else "LOW")
+
+        recs.append(Recommendation(
+            id=rec_id,
+            description=(
+                f"Upsample minority class {min_cls} from {min_count} to {target_count} rows "
+                f"via random oversampling. "
+                f"Imbalance ratio: {imbalance_ratio:.2f} → {ratio_after:.2f}. "
+                f"Class entropy: {entropy_pct:.1f}% of maximum. Risk: {risk}.{overbalance_note}"
+            ),
+            action_type="oversample",
+            estimated_impact=round(min(0.05 + (1 - imbalance_ratio) * 0.15, 0.12), 3),
+            confidence=round(0.60 if is_overbalance else 0.80, 2),
+            session_id=sid,
+            _payload={
+                "action": "oversample",
+                "class": min_cls,
+                "target_count": target_count,
+                "overbalance": is_overbalance,
+            },
+        ))
+        rec_id += 1
+
+        if imbalance_ratio < 0.5:
+            undersample_target = min_count * 2
+            recs.append(Recommendation(
+                id=rec_id,
+                description=(
+                    f"Downsample majority class {max_cls} from {max_count} to {undersample_target} rows "
+                    f"via random undersampling. "
+                    f"Warning: loses {max_count - undersample_target} majority-class examples. "
+                    f"Risk: MEDIUM — use only if dataset is large enough."
+                ),
+                action_type="undersample",
+                estimated_impact=round(min(0.03 + (1 - imbalance_ratio) * 0.08, 0.08), 3),
+                confidence=0.75,
+                session_id=sid,
+                _payload={
+                    "action": "undersample",
+                    "class": max_cls,
+                    "target_count": undersample_target,
+                },
+            ))
+
+        session_registry.register(recs)
+        return recs
+
+
+# ── ValidatorAgent ────────────────────────────────────────────────────────────
+
+class ValidatorAgent:
+    """
+    Checks working_copy against column metadata for rule violations.
+    Returns list of Violation objects (diagnostic only — not recommendations).
+    Costs 2 budget per call.
+    ~10% false positive rate (flagged with [FALSE POSITIVE WARNING]).
+    """
+
+    def query(self, df: pd.DataFrame, col_meta: Dict) -> List[Violation]:
+        rng = _seeded_rng(df, "validator")
+        violations = []
+
+        for col, meta in col_meta.items():
+            if col == "target" or col not in df.columns:
+                continue
+
+            expected_dtype = meta.get("expected_dtype", "float64")
+            valid_range = meta.get("valid_range")
+
+            # Type violations
+            if expected_dtype in ("float64", "int64"):
+                n_errors = sum(1 for val in df[col].dropna() if not _is_numeric(val))
+                if n_errors > 0:
+                    pct = n_errors / len(df) * 100
+                    violations.append(Violation(
+                        column=col,
+                        rule=f"dtype={expected_dtype}",
+                        count=n_errors,
+                        description=f"{n_errors} non-numeric values in '{col}' ({pct:.1f}%). Recommend fix_type_errors.",
+                        severity="CRITICAL" if pct > 10 else "WARNING",
+                    ))
+
+            # Range violations
+            if valid_range:
+                lo, hi = valid_range
+                try:
+                    numeric_vals = pd.to_numeric(df[col], errors="coerce").dropna()
+                    n_out = int(((numeric_vals < lo) | (numeric_vals > hi)).sum())
+                    if n_out > 0:
+                        max_val = float(numeric_vals.max())
+                        min_val = float(numeric_vals.min())
+                        std = float(numeric_vals.std()) or 1.0
+                        z_max = abs(max_val - numeric_vals.mean()) / std
+                        violations.append(Violation(
+                            column=col,
+                            rule=f"range=[{lo},{hi}]",
+                            count=n_out,
+                            description=(
+                                f"{n_out} values in '{col}' outside [{lo}, {hi}]. "
+                                f"Observed range: [{min_val:.1f}, {max_val:.1f}] "
+                                f"(max Z-score: {z_max:.1f}). "
+                                f"Severity: {'CRITICAL — likely data corruption' if z_max > 5 else 'WARNING — possible outliers'}."
+                            ),
+                            severity="CRITICAL" if z_max > 5 else "WARNING",
+                        ))
+                except Exception:
+                    pass
+
+        # ~10% false positive
+        fp_hash = int(hashlib.md5(df.to_json().encode()).hexdigest()[:4], 16)
+        if fp_hash % 100 < 10:
+            feature_cols = [c for c in df.columns if c != "target"]
+            if feature_cols:
+                fp_col = rng.choice(feature_cols)
+                violations.append(Violation(
+                    column=fp_col,
+                    rule="distribution_check",
+                    count=rng.randint(1, 5),
+                    description=(
+                        f"[FALSE POSITIVE WARNING] Unusual value distribution in '{fp_col}' "
+                        f"— may not be a real issue. Verify before acting."
+                    ),
+                    severity="INFO",
+                ))
+
+        return violations
+
+
+# ── AnalystAgent ──────────────────────────────────────────────────────────────
+
+class AnalystAgent:
+    """
+    Meta-specialist that performs holistic dataset diagnosis.
+    Returns a prioritised action plan rather than individual recommendations.
+    Costs 1 budget.
+
+    Analyses:
+    - Missing value severity
+    - Class imbalance severity
+    - Type error severity
+    - Remaining accuracy gap
+    Then ranks problems and recommends an ordered sequence of specialist calls.
+    """
+
+    def query(
+        self,
+        df: pd.DataFrame,
+        col_meta: Dict,
+        current_accuracy: float,
+        target_accuracy: float,
+        budget_remaining: int,
+    ) -> str:
+        """Return a formatted diagnostic + action plan string."""
+
+        n_rows = max(len(df), 1)
+        n_cells = n_rows * max(len(df.columns) - 1, 1)
+
+        # ── Score each problem dimension (0.0 – 1.0) ──────────────────────
+
+        # 1. Missing value severity
+        total_missing = int(df.isnull().sum().sum())
+        missing_severity = min(1.0, total_missing / n_cells * 5)
+
+        # 2. Class imbalance severity
+        vc = df["target"].value_counts()
+        if len(vc) >= 2:
+            imbalance_severity = 1.0 - (vc.min() / vc.max())
+        else:
+            imbalance_severity = 0.0
+
+        # 3. Type error severity
+        n_type_errors = 0
+        for col in df.columns:
+            if col == "target":
+                continue
+            meta = col_meta.get(col, {})
+            if meta.get("expected_dtype", "float64") in ("float64", "int64"):
+                n_type_errors += sum(1 for val in df[col].dropna() if not _is_numeric(val))
+        type_severity = min(1.0, n_type_errors / n_cells * 10)
+
+        # 4. Accuracy gap
+        accuracy_gap = max(0.0, target_accuracy - current_accuracy)
+
+        # ── Rank problems ─────────────────────────────────────────────────
+        problems = [
+            ("class imbalance",  imbalance_severity, "query_balancer"),
+            ("missing values",   missing_severity,   "query_cleaner"),
+            ("type errors",      type_severity,      "query_cleaner"),
+        ]
+        problems.sort(key=lambda x: -x[1])
+
+        # ── Build diagnosis section ───────────────────────────────────────
+        diagnosis_lines = ["DIAGNOSIS:"]
+        for name, severity, specialist in problems:
+            if severity < 0.05:
+                level = "NONE"
+            elif severity < 0.3:
+                level = "LOW"
+            elif severity < 0.6:
+                level = "MEDIUM"
+            else:
+                level = "HIGH"
+            diagnosis_lines.append(
+                f"  - {name.title()}: severity={severity:.2f} [{level}] -> use {specialist}"
+            )
+        diagnosis_lines.append(
+            f"  - Accuracy gap: {accuracy_gap:.4f} "
+            f"({'within reach' if accuracy_gap < 0.05 else 'significant gap'})"
+        )
+
+        # ── Build action plan ─���───────────────────────────────────────────
+        plan_lines = [f"\nRECOMMENDED PLAN (budget remaining: {budget_remaining}):"]
+        step = 1
+
+        # Recommend top 2 non-trivial problems
+        for name, severity, specialist in problems:
+            if severity >= 0.1:
+                plan_lines.append(f"  {step}. {specialist} → apply best recommendation")
+                step += 1
+            if step > 3:
+                break
+
+        # Always validate after fixes
+        plan_lines.append(f"  {step}. validate (check accuracy improvement)")
+        step += 1
+
+        # Budget guidance
+        if budget_remaining <= 8:
+            plan_lines.append(
+                f"  {step}. submit NOW — budget is critically low ({budget_remaining} steps left)"
+            )
+            plan_lines.append("  NOTE: Skip query_validator (costs 2 budget).")
+        elif accuracy_gap < 0.02:
+            plan_lines.append(f"  {step}. submit — you are very close to target")
+        else:
+            plan_lines.append(f"  {step}. Repeat if accuracy gap remains > 0.02, then submit")
+
+        # Feature note
+        if imbalance_severity > missing_severity and imbalance_severity > 0.2:
+            plan_lines.append("\nPRIORITY NOTE: Class imbalance is the dominant issue — fix this first.")
+        elif missing_severity > 0.2:
+            plan_lines.append("\nPRIORITY NOTE: High missing-value rate — clean data before augmenting.")
+
+        return "\n".join(diagnosis_lines + plan_lines)

sft_generator.py

@@ -0,0 +1,359 @@
+"""
+Ultra-Fast SFT Data Generator — No sklearn, No Environment Execution.
+
+Instead of running the environment live, we generate realistic prompt/response
+pairs directly from templates using known dataset states.
+
+This is correct because:
+ - We know exactly what inspect_dataset returns (from dataset_generator)
+ - We know what query_cleaner returns (from specialist_agents)
+ - We know the reward trajectory
+ - The actual RL training will run the real environment — SFT just warms up
+   the LLM's action distribution (command grammar + strategy)
+
+Output: ~1000+ diverse examples in under 10 seconds.
+"""
+
+import json
+import os
+import random
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from server.dataset_generator import TASK_CONFIGS
+
+rng = random.Random(42)
+
+TASKS = list(TASK_CONFIGS.keys())
+
+# ── Prompt templates ──────────────────────────────────────────────────────────
+
+def make_prompt(
+    task: str,
+    step: int,
+    max_steps: int,
+    current_acc: float,
+    target_acc: float,
+    baseline_acc: float,
+    dataset_shape: str,
+    rows_pct: float,
+    quality: float,
+    budget: int,
+    session: str,
+    validate_left: int,
+    last_obs: str,
+) -> str:
+    gap = max(0.0, target_acc - current_acc)
+    return (
+        f"You are a Data-Centric AI agent improving an ML dataset.\n\n"
+        f"Task: {task}\n"
+        f"Step: {step}/{max_steps}\n"
+        f"Current accuracy: {current_acc:.4f}  "
+        f"Target: {target_acc:.4f}  Gap: {gap:.4f}\n"
+        f"Baseline accuracy: {baseline_acc:.4f}\n"
+        f"Dataset: {dataset_shape} | "
+        f"Rows preserved: {rows_pct*100:.1f}%\n"
+        f"Quality score: {quality:.4f} | "
+        f"Budget remaining: {budget}\n"
+        f"Active session: {session} | "
+        f"Validate calls left: {validate_left}\n\n"
+        f"Last observation:\n{last_obs}\n\n"
+        f"What is your next command?"
+    )
+
+
+# ── Observation text snippets ────────────────────────────────────────────────
+
+INSPECT_OBS_TEMPLATES = [
+    "=== Dataset Inspection ===\nShape: {rows} rows × {cols} features\nOriginal rows: {rows} | Preserved: 100.0%\nDuplicates: {dups}\nMissing values:\n  {col}: {missing}\nClass distribution: {dist}\nDtypes: {{'age': 'float64', 'score': 'float64', 'target': 'int64'}}",
+    "=== Dataset Inspection ===\nShape: {rows} rows × {cols} features\nDuplicates: {dups}\nMissing values:\n  {col}: {missing}\nClass distribution: {dist}",
+]
+
+INSPECT_MODEL_TEMPLATES = [
+    "=== Model Inspection ===\nAccuracy: {acc:.4f}\n  Class 0: precision={p0:.3f} recall={r0:.3f} f1={f0:.3f}\n  Class 1: precision={p1:.3f} recall={r1:.3f} f1={f1:.3f}\nTarget: {target:.4f} | Not yet",
+    "=== Model Inspection (cached) ===\nAccuracy: {acc:.4f}\nTarget: {target:.4f} | Not yet",
+]
+
+CLEANER_OBS_TEMPLATES = [
+    "=== Cleaner Recommendations ===\n[1] Fill {n} missing values in '{col}' using mean ({mean:.2f})\n    type=fill_missing impact=+0.075 confidence=0.90\n[2] Remove {dups} duplicate rows\n    type=remove_duplicates impact=+0.020 confidence=0.95",
+    "=== Cleaner Recommendations ===\n[1] Fill {n} missing values in '{col}' using mean ({mean:.2f})\n    type=fill_missing impact=+0.075 confidence=0.90\n[2] Fix {typos} type errors in 'income'\n    type=fix_type_errors impact=+0.040 confidence=0.75",
+    "=== Cleaner Recommendations ===\n[1] Fill {n} missing values in '{col}' using mean ({mean:.2f})\n    type=fill_missing impact=+0.075 confidence=0.90",
+]
+
+BALANCER_OBS_TEMPLATES = [
+    "=== Balancer Recommendations ===\n[1] Upsample minority class 1 from {min_c} to {maj_c} rows via random oversampling (imbalance ratio: {ratio:.2f})\n    type=oversample impact=+0.053 confidence=0.80",
+    "=== Balancer Recommendations ===\n[1] Downsample majority class 0 from {maj_c} to {min_c} rows\n    type=undersample impact=+0.030 confidence=0.70",
+]
+
+APPLY_OBS_TEMPLATES = [
+    "Applied: fill_missing [Fill {n} missing values in '{col}' using mean ({mean:.2f})]\n\nDataset health check:\n  Missing values: {remaining} remaining (was {was})\n  Duplicates: ✓ (was 0)\n  Row count: {rows}/{orig} (100.0% preserved)\n\nEstimated quality score: {quality:.4f}\nBudget remaining: {budget}",
+    "Applied: remove_duplicates [Remove {dups} duplicate rows]\n\nDataset health check:\n  Missing values: {remaining} remaining (was {was})\n  Duplicates: ✓ (was {dups})\n  Row count: {rows}/{orig} ({pct:.1f}% preserved)\n\nEstimated quality score: {quality:.4f}\nBudget remaining: {budget}",
+    "Applied: oversample [Upsample minority class 1 via random oversampling]\n\nDataset health check:\n  Missing values: 0 remaining (was 0)\n  Duplicates: 2 remaining (was 0)\n  Row count: {rows}/{orig} (102.0% preserved)\n\nEstimated quality score: {quality:.4f}\nBudget remaining: {budget}",
+]
+
+VALIDATE_OBS_TEMPLATES = [
+    "=== Validate ===\nRF Accuracy: {acc:.4f}  (primary)\nLR Accuracy: {lr_acc:.4f}  (secondary)\nAgreement:   BOTH_AGREE_IMPROVE -- fix is robust and generalises\n  Class 0: p={p:.3f} r={r:.3f} f1={f:.3f}\n  Class 1: p={p:.3f} r={r:.3f} f1={f:.3f}\nTarget: {target:.4f} | {status}",
+    "=== Validate ===\nRF Accuracy: {acc:.4f}  (primary)\nLR Accuracy: {lr_acc:.4f}  (secondary)\nAgreement:   BOTH_AGREE_IMPROVE -- fix is robust and generalises\nTarget: {target:.4f} | {status}",
+    "=== Validate (cached) ===\nRF Accuracy: {acc:.4f}  (primary)\nLR Accuracy: {lr_acc:.4f}  (secondary)\nTarget: {target:.4f} | {status}",
+]
+
+ERROR_OBS_TEMPLATES = [
+    "Error: Recommendation 1 has already been applied this session. Duplicate apply not allowed.",
+    "Validate on cooldown. Take 1 more action(s) before validating again.",
+    "Error: stale recommendation ID 99. Please re-query for fresh recommendations.",
+]
+
+RESET_OBS = (
+    "Episode started: {task}\n"
+    "Baseline accuracy: {baseline:.4f} | Target: {target:.4f}\n"
+    "Dataset: {rows} rows x {cols} features\n"
+    "Budget: {budget} steps\n\n"
+    "Available commands:\n"
+    "  inspect_dataset          - shape, dtypes, missing, class distribution\n"
+    "  inspect_model            - accuracy (RF + LR), F1, feature importance\n"
+    "  query_analyst            - holistic diagnosis + prioritised action plan (costs 1 budget)\n"
+    "  query_cleaner            - get cleaning recommendations\n"
+    "  query_augmenter [class]  - get augmentation suggestions\n"
+    "  query_balancer           - get resampling recommendations\n"
+    "  query_validator          - check rule violations (costs 2 budget)\n"
+    "  apply [id]               - apply recommendation by ID\n"
+    "  reject [id]              - reject a recommendation\n"
+    "  validate                 - retrain and score (cooldown applies)\n"
+    "  submit                   - finalize episode"
+)
+
+ANALYST_OBS_TEMPLATES = [
+    "=== Analyst Report (costs 1 budget) ===\nDIAGNOSIS:\n  - Class Imbalance: severity={imb:.2f} [HIGH] -> use query_balancer\n  - Missing Values: severity={miss:.2f} [MEDIUM] -> use query_cleaner\n  - Type Errors: severity=0.00 [NONE]\n  - Accuracy gap: {gap:.4f} (significant gap)\n\nRECOMMENDED PLAN (budget remaining: {budget}):\n  1. query_balancer -> apply best recommendation\n  2. query_cleaner -> apply best recommendation\n  3. validate (check accuracy improvement)\n  4. submit if accuracy >= target\n\nPRIORITY NOTE: Class imbalance is the dominant issue -- fix this first.",
+    "=== Analyst Report (costs 1 budget) ===\nDIAGNOSIS:\n  - Missing Values: severity={miss:.2f} [HIGH] -> use query_cleaner\n  - Class Imbalance: severity={imb:.2f} [LOW] -> use query_balancer\n  - Type Errors: severity=0.00 [NONE]\n  - Accuracy gap: {gap:.4f} (significant gap)\n\nRECOMMENDED PLAN (budget remaining: {budget}):\n  1. query_cleaner -> apply best recommendation\n  2. query_balancer -> apply best recommendation\n  3. validate\n  4. submit",
+]
+
+
+# ── Episode builders ─────────────────────────────────────────────────────────
+
+def sample_dataset_params(task: str, seed: int):
+    """Sample realistic dataset params for a given task."""
+    cfg = TASK_CONFIGS[task]
+    rng2 = random.Random(seed)
+    rows_map = {"task_0_tutorial": 100, "task_1_easy": 200,
+                "task_2_medium": 500, "task_3_hard": 900}
+    cols_map = {"task_0_tutorial": 4, "task_1_easy": 5,
+                "task_2_medium": 7, "task_3_hard": 10}
+    rows = rows_map[task]
+    cols = cols_map[task]
+    missing_cols = ["age", "income", "score"][:rng2.randint(1, 3)]
+    missing_pct = rng2.uniform(0.10, 0.30)
+    n_missing = int(rows * missing_pct)
+    mean_val = rng2.uniform(30.0, 60.0)
+    dups = rng2.randint(0, int(rows * 0.05))
+    maj_class = int(rows * rng2.uniform(0.52, 0.65))
+    min_class = rows - maj_class
+    return {
+        "task": task, "rows": rows, "cols": cols,
+        "missing_col": missing_cols[0], "n_missing": n_missing,
+        "mean_val": round(mean_val, 2), "dups": dups,
+        "maj_class": maj_class, "min_class": min_class,
+        "baseline": cfg["baseline_accuracy"],
+        "target": cfg["target_accuracy"],
+        "budget": cfg["budget"],
+    }
+
+
+def build_episode(task: str, seed: int, strategy: list) -> list:
+    """
+    Build a synthetic SFT episode using template obs + fixed action sequence.
+    Returns list of {prompt, response} dicts.
+    """
+    p = sample_dataset_params(task, seed)
+    cfg = TASK_CONFIGS[task]
+    examples = []
+
+    acc = p["baseline"]
+    quality = round(rng.uniform(0.45, 0.65), 4)
+    rows = p["rows"]
+    missing_remaining = p["n_missing"]
+    budget = p["budget"]
+    session = "none"
+    validate_left = 3
+    prev_obs = RESET_OBS.format(
+        task=task, baseline=p["baseline"], target=p["target"],
+        rows=rows, cols=p["cols"], budget=budget
+    )
+
+    for step, action in enumerate(strategy):
+        prompt = make_prompt(
+            task=task, step=step, max_steps=p["budget"],
+            current_acc=acc, target_acc=p["target"], baseline_acc=p["baseline"],
+            dataset_shape=f"{rows} rows × {p['cols']} columns",
+            rows_pct=rows / p["rows"], quality=quality, budget=budget,
+            session=session, validate_left=validate_left, last_obs=prev_obs,
+        )
+        examples.append({"prompt": prompt, "response": action})
+
+        # Simulate observation update
+        budget -= 1
+        cmd = action.split()[0].lower()
+
+        if cmd == "inspect_dataset":
+            t = rng.choice(INSPECT_OBS_TEMPLATES)
+            dist = f"class 0: {p['maj_class']}, class 1: {p['min_class']}"
+            prev_obs = t.format(
+                rows=rows, cols=p["cols"], dups=p["dups"],
+                col=p["missing_col"], missing=missing_remaining, dist=dist,
+            )
+        elif cmd == "inspect_model":
+            t = rng.choice(INSPECT_MODEL_TEMPLATES)
+            p0 = round(rng.uniform(0.55, 0.75), 3)
+            r0 = round(rng.uniform(0.55, 0.75), 3)
+            prev_obs = t.format(
+                acc=acc, target=p["target"],
+                p0=p0, r0=r0, f0=round(2*p0*r0/(p0+r0+1e-9), 3),
+                p1=p0, r1=r0, f1=round(2*p0*r0/(p0+r0+1e-9), 3),
+            )
+        elif cmd == "query_cleaner":
+            t = rng.choice(CLEANER_OBS_TEMPLATES)
+            session = f"cleaner:{seed:08x}"
+            prev_obs = t.format(
+                n=missing_remaining, col=p["missing_col"],
+                mean=p["mean_val"], dups=p["dups"], typos=rng.randint(2, 8),
+            )
+        elif cmd == "query_balancer":
+            t = rng.choice(BALANCER_OBS_TEMPLATES)
+            session = f"balancer:{seed:08x}"
+            ratio = round(p["min_class"] / max(p["maj_class"], 1), 2)
+            prev_obs = t.format(
+                min_c=p["min_class"], maj_c=p["maj_class"], ratio=ratio
+            )
+        elif cmd == "query_augmenter":
+            session = f"augmenter:{seed:08x}"
+            cls = action.split()[1] if len(action.split()) > 1 else "0"
+            n_synth = rng.randint(5, 25)
+            prev_obs = (
+                f"=== Augmenter Recommendations ===\n"
+                f"[1] Synthesize {n_synth} samples for class {cls} via SMOTE\n"
+                f"    type=augment_class impact=+0.040 confidence=0.72"
+            )
+        elif cmd == "query_analyst":
+            budget -= 1  # costs 1 extra
+            t = rng.choice(ANALYST_OBS_TEMPLATES)
+            imb = round(rng.uniform(0.3, 0.8), 2)
+            miss = round(rng.uniform(0.1, 0.5), 2)
+            gap = round(p["target"] - acc, 4)
+            prev_obs = t.format(imb=imb, miss=miss, gap=gap, budget=budget)
+        elif cmd == "query_validator":
+            budget -= 1  # costs 2
+            prev_obs = (
+                "=== Validator Report (costs 2 budget) ===\n"
+                f"  [WARNING] [{p['missing_col']}] rule=no_missing "
+                f"count={missing_remaining}\n"
+                f"    Column '{p['missing_col']}' has {missing_remaining} missing values."
+            )
+        elif cmd == "apply":
+            rec_id = int(action.split()[1]) if len(action.split()) > 1 else 1
+            t = rng.choice(APPLY_OBS_TEMPLATES)
+            was_missing = missing_remaining
+            missing_remaining = max(0, missing_remaining - p["n_missing"])
+            quality = min(1.0, quality + rng.uniform(0.10, 0.35))
+            quality = round(quality, 4)
+            prev_obs = t.format(
+                n=p["n_missing"], col=p["missing_col"], mean=p["mean_val"],
+                remaining=missing_remaining, was=was_missing,
+                rows=rows, orig=p["rows"], pct=rows/p["rows"]*100,
+                dups=p["dups"], quality=quality, budget=budget,
+            )
+        elif cmd == "reject":
+            prev_obs = f"Recommendation {action.split()[1] if len(action.split())>1 else 1} rejected."
+        elif cmd == "validate":
+            if validate_left > 0:
+                acc = min(1.0, acc + rng.uniform(0.05, 0.35))
+                acc = round(acc, 4)
+                lr_acc = round(min(1.0, acc + rng.uniform(-0.03, 0.03)), 4)
+                validate_left -= 1
+                t = rng.choice(VALIDATE_OBS_TEMPLATES)
+                status = "HIT v" if acc >= p["target"] else "Not yet"
+                pv = round(rng.uniform(0.75, 0.98), 3)
+                rv = round(rng.uniform(0.75, 0.98), 3)
+                prev_obs = t.format(
+                    acc=acc, lr_acc=lr_acc, target=p["target"], status=status,
+                    p=pv, r=rv, f=round(2*pv*rv/(pv+rv+1e-9), 3),
+                )
+            else:
+                prev_obs = "Validate on cooldown. Take 2 more action(s) before validating again."
+        elif cmd == "submit":
+            break
+
+    return examples
+
+
+# ── Strategy sequences ────────────────────────────────────────────────────────
+
+STRATEGIES = {
+    "minimal_clean":             ["inspect_dataset", "query_cleaner", "apply 1", "apply 2", "inspect_dataset", "validate", "submit"],
+    "inspect_model_first":       ["inspect_dataset", "inspect_model", "query_cleaner", "apply 1", "inspect_dataset", "validate", "submit"],
+    "clean_then_balance":        ["inspect_dataset", "query_cleaner", "apply 1", "apply 2", "query_balancer", "apply 1", "inspect_dataset", "validate", "submit"],
+    "reject_then_apply":         ["inspect_dataset", "query_cleaner", "reject 1", "apply 2", "inspect_dataset", "validate", "submit"],
+    "baseline_validate_first":   ["inspect_dataset", "validate", "query_cleaner", "apply 1", "inspect_dataset", "validate", "submit"],
+    "augment_path":              ["inspect_dataset", "query_cleaner", "apply 1", "query_augmenter 0", "apply 1", "inspect_dataset", "validate", "submit"],
+    "with_validator":            ["inspect_dataset", "query_validator", "query_cleaner", "apply 1", "inspect_dataset", "validate", "submit"],
+    "deep_clean_requery":        ["inspect_dataset", "query_cleaner", "apply 1", "apply 2", "query_cleaner", "apply 1", "inspect_dataset", "validate", "submit"],
+    "fast_submit":               ["query_cleaner", "apply 1", "apply 2", "inspect_dataset", "submit"],
+    "balance_heavy":             ["inspect_dataset", "query_balancer", "apply 1", "query_cleaner", "apply 1", "inspect_dataset", "validate", "submit"],
+    "reject_requery":            ["inspect_dataset", "query_cleaner", "reject 1", "reject 2", "query_cleaner", "apply 1", "inspect_dataset", "validate", "submit"],
+    "multi_augment":             ["inspect_dataset", "query_cleaner", "apply 1", "query_augmenter 1", "apply 1", "inspect_dataset", "validate", "submit"],
+    "model_then_balance":        ["inspect_model", "query_balancer", "apply 1", "inspect_dataset", "validate", "submit"],
+    "full_pipeline":             ["inspect_dataset", "inspect_model", "query_cleaner", "apply 1", "query_balancer", "apply 1", "query_augmenter 0", "apply 1", "inspect_dataset", "validate", "submit"],
+    "suboptimal_no_validate":    ["inspect_dataset", "query_cleaner", "apply 1", "submit"],
+    "inspect_only_submit":       ["inspect_dataset", "inspect_model", "submit"],
+    "reject_all_then_requery":   ["inspect_dataset", "query_cleaner", "reject 1", "reject 2", "query_balancer", "apply 1", "inspect_dataset", "validate", "submit"],
+    "apply3_then_validate":      ["inspect_dataset", "query_cleaner", "apply 1", "apply 2", "query_balancer", "apply 1", "query_augmenter 0", "apply 1", "inspect_dataset", "validate", "submit"],
+    # NEW: analyst-led strategies
+    "analyst_led_clean":         ["query_analyst", "inspect_dataset", "query_cleaner", "apply 1", "apply 2", "validate", "submit"],
+    "analyst_led_balance":       ["query_analyst", "query_balancer", "apply 1", "query_cleaner", "apply 1", "validate", "submit"],
+    "analyst_full_pipeline":     ["query_analyst", "inspect_dataset", "inspect_model", "query_cleaner", "apply 1", "query_balancer", "apply 1", "validate", "submit"],
+}
+
+
+def generate_sft_data(output_file: str = "sft_data.jsonl", seeds_per_combo: int = 15):
+    sft_examples = []
+
+    print(f"Generating SFT data: {len(STRATEGIES)} strategies × {len(TASKS)} tasks × {seeds_per_combo} seeds")
+
+    for strategy_name, sequence in STRATEGIES.items():
+        strategy_examples = []
+        for task in TASKS:
+            for seed in range(seeds_per_combo):
+                episode = build_episode(task, seed, sequence)
+                strategy_examples.extend(episode)
+        sft_examples.extend(strategy_examples)
+        print(f"  {strategy_name:<30} +{len(strategy_examples)} examples")
+
+    rng.shuffle(sft_examples)
+
+    out_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), output_file)
+    with open(out_path, "w", encoding="utf-8") as f:
+        for ex in sft_examples:
+            f.write(json.dumps(ex) + "\n")
+
+    # Diversity report
+    from collections import Counter
+    responses = [ex["response"] for ex in sft_examples]
+    unique_cmds = set(responses)
+
+    print(f"\n{'='*55}")
+    print(f"Total examples:   {len(sft_examples)}")
+    print(f"Unique commands:  {len(unique_cmds)}")
+    print(f"Unique prompts:   {len(set(ex['prompt'] for ex in sft_examples))}")
+    print(f"\nResponse distribution:")
+    for cmd, cnt in Counter(responses).most_common():
+        pct = cnt / len(responses) * 100
+        bar = "#" * int(pct / 2)
+        flag = " ← DOMINANT" if pct > 25 else ""
+        print(f"  {cmd:<32} {cnt:>5}  ({pct:5.1f}%)  {bar}{flag}")
+
+    print(f"\nOutput: {out_path}")
+    print("✓ SFT generation complete (no sklearn, instant).")
+    return sft_examples
+
+
+if __name__ == "__main__":
+    generate_sft_data()

test_features_smoke.py

@@ -0,0 +1,80 @@
+# -*- coding: utf-8 -*-
+"""Smoke test: verify all 5 new features work end-to-end."""
+import sys, os
+sys.stdout.reconfigure(encoding='utf-8', errors='replace')
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from models import DataCentricAction
+from server.data_centric_environment import DataCentricEnvironment
+
+env = DataCentricEnvironment()
+obs = env.reset(task='task_1_easy', seed=7)
+print(f"Reset OK. Budget: {obs.budget_remaining}, Baseline: {obs.baseline_accuracy:.4f}")
+print()
+
+# ─── Feature 3: query_analyst ────────────────────────────────────────────────
+print("=" * 60)
+print("TEST 1: query_analyst (meta-specialist, costs 1 budget)")
+print("=" * 60)
+budget_before = obs.budget_remaining
+obs = env.step(DataCentricAction(message='query_analyst'))
+budget_after = obs.budget_remaining
+print(obs.response)
+print(f"\n[BUDGET CHECK] Before={budget_before}, After={budget_after}, Diff={budget_before - budget_after}")
+assert "DIAGNOSIS" in obs.response, "FAIL: no DIAGNOSIS section"
+assert "RECOMMENDED PLAN" in obs.response, "FAIL: no RECOMMENDED PLAN section"
+assert budget_before - budget_after == 2, f"FAIL: should cost 2 total (1 cmd + 1 analyst), got {budget_before - budget_after}"
+print("PASS: query_analyst works")
+
+# ─── Feature 1: Smarter specialists ─────────────────────────────────────────
+print()
+print("=" * 60)
+print("TEST 2: query_cleaner (smarter specialists with reasoning)")
+print("=" * 60)
+obs = env.step(DataCentricAction(message='query_cleaner'))
+print(obs.response)
+# Check for statistical reasoning markers
+has_reasoning = any(kw in obs.response for kw in ["skew", "Risk:", "Reason:", "median", "mean", "%"])
+assert has_reasoning, "FAIL: no statistical reasoning found in cleaner output"
+print("PASS: smarter specialists working (statistical reasoning present)")
+
+# ─── Feature 5: Drift detection ──────────────────────────────────────────────
+print()
+print("=" * 60)
+print("TEST 3: apply 1 (drift detection after apply)")
+print("=" * 60)
+obs = env.step(DataCentricAction(message='apply 1'))
+print(obs.response)
+has_drift = "Distribution drift" in obs.response or "drift" in obs.response.lower()
+assert has_drift, "FAIL: no drift information in apply response"
+print("PASS: drift detection working")
+
+# ─── Feature 2 + 4: Dual classifier + Feature importance ───────────────────
+print()
+print("=" * 60)
+print("TEST 4: validate (dual classifier + feature importance)")
+print("=" * 60)
+obs = env.step(DataCentricAction(message='validate'))
+print(obs.response)
+assert "RF Accuracy" in obs.response, "FAIL: no RF Accuracy"
+assert "LR Accuracy" in obs.response, "FAIL: no LR Accuracy"
+assert "Agreement" in obs.response, "FAIL: no Agreement signal"
+has_feat_imp = "Feature importance" in obs.response
+print(f"Feature importance shown: {has_feat_imp}")
+print("PASS: dual classifier + agreement signal working")
+
+# ─── Feature 4: Feature importance in inspect_model ─────────────────────────
+print()
+print("=" * 60)
+print("TEST 5: inspect_model (RF + LR + feature importance)")
+print("=" * 60)
+obs = env.step(DataCentricAction(message='inspect_model'))
+print(obs.response)
+assert "RF Accuracy" in obs.response, "FAIL: no RF Accuracy in inspect_model"
+assert "LR Accuracy" in obs.response, "FAIL: no LR Accuracy in inspect_model"
+print("PASS: inspect_model shows dual classifier")
+
+print()
+print("=" * 60)
+print("ALL 5 FEATURES VERIFIED OK")
+print("=" * 60)

tests/conftest.py

@@ -0,0 +1,2 @@
+import sys, os
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

tests/test_environment.py

@@ -0,0 +1,158 @@
+"""
+Unit tests for anti-exploit and environment stability.
+
+Tests that the core safety invariants hold:
+- ground truth never mutates
+- budget is enforced
+- validate calls are limited
+- undo works correctly
+
+Run with: pytest tests/test_environment.py -v
+"""
+import pytest
+import sys, os
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from models import DataCentricAction
+from server.data_centric_environment import DataCentricEnvironment
+
+
+# ── Helpers ───────────────────────────────────────────────────────────────────
+
+def make_env(task="task_0_tutorial", seed=42) -> tuple:
+    """Return (env, reset_obs)."""
+    env = DataCentricEnvironment()
+    obs = env.reset(task=task, seed=seed)
+    return env, obs
+
+
+def step(env, cmd: str):
+    return env.step(DataCentricAction(message=cmd))
+
+
+# ── Ground truth immutability ─────────────────────────────────────────────────
+
+class TestGroundTruth:
+
+    def test_ground_truth_unchanged_after_reset(self):
+        env, _ = make_env()
+        gt_before = env._ground_truth.copy()
+        env.reset(task="task_0_tutorial", seed=123)
+        # After re-reset, a NEW ground truth is loaded — old one doesn't matter
+        assert env._ground_truth is not None
+
+    def test_ground_truth_unchanged_after_query(self):
+        env, _ = make_env()
+        gt_before = env._ground_truth.copy()
+        step(env, "query_cleaner")
+        assert env._ground_truth.equals(gt_before), "GT mutated after query_cleaner"
+
+    def test_ground_truth_unchanged_after_inspect(self):
+        env, _ = make_env()
+        gt_before = env._ground_truth.copy()
+        step(env, "inspect_dataset")
+        assert env._ground_truth.equals(gt_before), "GT mutated after inspect_dataset"
+
+
+# ── Budget enforcement ────────────────────────────────────────────────────────
+
+class TestBudget:
+
+    def test_budget_decreases_each_step(self):
+        env, obs = make_env()
+        budget_start = obs.budget_remaining
+        obs2 = step(env, "inspect_dataset")
+        assert obs2.budget_remaining < budget_start
+
+    def test_done_after_budget_exhausted(self):
+        env, obs = make_env()
+        budget = obs.budget_remaining
+        last_obs = obs
+        for _ in range(budget + 5):
+            if last_obs.done:
+                break
+            last_obs = step(env, "inspect_dataset")
+        assert last_obs.done, "Episode should be done after budget exhausted"
+
+
+# ── Validate calls ────────────────────────────────────────────────────────────
+
+class TestValidateCalls:
+
+    def test_validate_calls_start_at_3(self):
+        env, obs = make_env()
+        assert obs.validate_calls_remaining == 3
+
+    def test_validate_call_decrements(self):
+        env, obs = make_env()
+        obs2 = step(env, "validate")
+        # Either decrement or cooldown message — either way calls consumed
+        assert obs2.validate_calls_remaining <= 3
+
+
+# ── Undo ─────────────────────────────────────────────────────────────────────
+
+class TestUndo:
+
+    def test_undo_without_history_returns_response(self):
+        """Undo with no history should return a message, not crash."""
+        env, _ = make_env()
+        obs = step(env, "undo")
+        assert obs.response is not None
+        assert len(obs.response) > 0
+
+    def test_undo_after_apply_restores_state(self):
+        """After apply+undo, working copy should match pre-apply state."""
+        env, _ = make_env()
+        step(env, "query_cleaner")
+        wc_before_apply = env._working_copy.copy()
+        step(env, "apply 1")
+        step(env, "undo")
+        # After undo, working copy should be restored
+        assert env._working_copy.shape == wc_before_apply.shape, (
+            f"Shape mismatch after undo: {env._working_copy.shape} vs {wc_before_apply.shape}"
+        )
+
+
+# ── Snapshot stack ────────────────────────────────────────────────────────────
+
+class TestSnapshotStack:
+
+    def test_max_3_snapshots(self):
+        env, _ = make_env()
+        step(env, "query_cleaner")
+        # Apply 4 times — stack should cap at 3
+        for i in range(1, 5):
+            step(env, f"apply {i}")
+            step(env, "query_cleaner")  # re-query for fresh recs
+        assert len(env._dataset_history) <= env._max_history
+
+    def test_snapshot_cleared_on_reset(self):
+        env, _ = make_env()
+        step(env, "query_cleaner")
+        step(env, "apply 1")
+        env.reset(task="task_0_tutorial", seed=99)
+        assert len(env._dataset_history) == 0
+
+
+# ── Reward sanity ─────────────────────────────────────────────────────────────
+
+class TestRewardSanity:
+
+    def test_unknown_command_gives_penalty(self):
+        env, _ = make_env()
+        obs = step(env, "blorp_invalid_command_xyz")
+        assert obs.reward <= 0.0, (
+            f"Unknown command should not give positive reward, got {obs.reward}"
+        )
+
+    def test_reward_within_range(self):
+        env, _ = make_env()
+        from server.grader import REWARD_MIN, REWARD_MAX
+        for cmd in ["inspect_dataset", "inspect_model", "query_cleaner",
+                    "query_balancer", "validate"]:
+            obs = step(env, cmd)
+            r = obs.reward
+            assert REWARD_MIN - 0.01 <= r <= REWARD_MAX + 0.01, (
+                f"Command '{cmd}' gave out-of-range reward {r}"
+            )

tests/test_grader.py

@@ -0,0 +1,195 @@
+"""
+Unit tests for server/grader.py
+
+Tests every reward component individually plus the range clamp.
+Run with: pytest tests/test_grader.py -v
+"""
+import importlib.util
+import sys
+import os
+import pytest
+import pandas as pd
+
+# Import grader directly (avoids server/__init__.py → openenv-core chain)
+_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+_GRADER = os.path.join(_ROOT, "server", "grader.py")
+spec = importlib.util.spec_from_file_location("grader", _GRADER)
+_mod = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(_mod)
+
+REWARD_MAX = _mod.REWARD_MAX
+REWARD_MIN = _mod.REWARD_MIN
+compute_accuracy_reward = _mod.compute_accuracy_reward
+compute_efficiency_reward = _mod.compute_efficiency_reward
+compute_preservation_reward = _mod.compute_preservation_reward
+compute_process_reward = _mod.compute_process_reward
+compute_step_reward = _mod.compute_step_reward
+compute_total_reward = _mod.compute_total_reward
+compute_lightweight_score = _mod.compute_lightweight_score
+
+
+# ── Accuracy reward ───────────────────────────────────────────────────────────
+
+class TestAccuracyReward:
+
+    def test_improvement_positive(self):
+        r = compute_accuracy_reward(0.70, 0.62, 0.62, 0.80)
+        assert r > 0, f"Improvement should give positive reward, got {r}"
+
+    def test_regression_negative(self):
+        r = compute_accuracy_reward(0.60, 0.70, 0.62, 0.80)
+        assert r < 0, f"Regression should give negative reward, got {r}"
+
+    def test_no_change_zero(self):
+        r = compute_accuracy_reward(0.65, 0.65, 0.62, 0.80)
+        assert r == 0.0
+
+    def test_submit_success_bonus(self):
+        r = compute_accuracy_reward(0.80, 0.75, 0.62, 0.80, is_submit=True)
+        assert r > 0.5, f"Submit success should add bonus, got {r}"
+
+    def test_submit_fail_partial_credit(self):
+        """Halfway to target should give some credit."""
+        # baseline=0.62, target=0.80, current=0.71 → 50% of the way there
+        r = compute_accuracy_reward(0.71, 0.70, 0.62, 0.80, is_submit=True)
+        # Current is at 0.71, previous was 0.70 → small improvement, partial credit
+        assert r > 0, f"Partial progress at submit should give credit, got {r}"
+
+
+# ── Preservation reward ───────────────────────────────────────────────────────
+
+class TestPreservationReward:
+
+    def test_above_90_bonus(self):
+        r = compute_preservation_reward(97, 100)
+        assert r > 0
+
+    def test_below_90_zero_or_neg(self):
+        r = compute_preservation_reward(85, 100)
+        assert r <= 0.02  # at best neutral at 85%
+
+    def test_below_50_catastrophic(self):
+        r = compute_preservation_reward(40, 100)
+        assert r <= -0.40, f"Expected catastrophic penalty, got {r}"
+
+    def test_full_preservation(self):
+        r = compute_preservation_reward(100, 100)
+        assert r == 0.05
+
+
+# ── Process reward ────────────────────────────────────────────────────────────
+
+class TestProcessReward:
+
+    def test_query_after_inspect_rewarded(self):
+        history = ["inspect_dataset"]
+        r = compute_process_reward(history, "query_cleaner")
+        assert r > 0
+
+    def test_apply_without_query_penalized(self):
+        history = ["inspect_dataset", "inspect_model"]
+        r = compute_process_reward(history, "apply 1")
+        assert r < 0
+
+    def test_apply_after_query_rewarded(self):
+        history = ["inspect_dataset", "query_cleaner"]
+        r = compute_process_reward(history, "apply 1")
+        assert r > 0
+
+    def test_submit_without_validate_penalized(self):
+        history = ["inspect_dataset", "query_cleaner", "apply 1"]
+        r = compute_process_reward(history, "submit")
+        assert r < 0
+
+
+# ── Step reward ───────────────────────────────────────────────────────────────
+
+class TestStepReward:
+
+    def test_quality_improvement_positive(self):
+        r = compute_step_reward("apply 1", quality_before=0.5, quality_after=0.7,
+                                rows_preserved_after=0.97)
+        assert r > 0
+
+    def test_quality_degradation_negative(self):
+        r = compute_step_reward("apply 1", quality_before=0.7, quality_after=0.4,
+                                rows_preserved_after=0.97)
+        assert r < 0
+
+    def test_non_apply_zero(self):
+        r = compute_step_reward("validate", quality_before=0.5, quality_after=0.7,
+                                rows_preserved_after=0.97)
+        assert r == 0.0
+
+    def test_low_preservation_penalty(self):
+        r = compute_step_reward("apply 1", quality_before=0.5, quality_after=0.6,
+                                rows_preserved_after=0.75)
+        # Row preservation penalty should reduce reward
+        r_without = compute_step_reward("apply 1", quality_before=0.5, quality_after=0.6,
+                                        rows_preserved_after=0.97)
+        assert r < r_without
+
+
+# ── Total reward range ────────────────────────────────────────────────────────
+
+class TestRewardRange:
+
+    def test_within_declared_range(self):
+        """All reward combinations must stay within [-1.0, 1.0]."""
+        test_cases = [
+            (0.5, 0.1, 0.05, 0.2, 0.1),
+            (-0.8, -0.1, -0.4, -0.05, -0.2),
+            (1.0, 0.2, 0.05, 0.2, 0.15),   # might need clamping
+            (-1.5, -0.2, -0.4, -0.05, -0.3),  # definitely needs clamping
+        ]
+        for acc, proc, pres, eff, step in test_cases:
+            r = compute_total_reward(acc, proc, pres, eff, step)
+            assert REWARD_MIN <= r <= REWARD_MAX, (
+                f"Reward {r} out of [{REWARD_MIN}, {REWARD_MAX}] "
+                f"for inputs acc={acc} proc={proc} pres={pres}"
+            )
+
+    def test_clamping_applied(self):
+        """Extreme inputs should be clamped, not crash."""
+        r = compute_total_reward(10.0, 5.0, 5.0)
+        assert r == REWARD_MAX
+
+        r = compute_total_reward(-10.0, -5.0, -5.0)
+        assert r == REWARD_MIN
+
+
+# ── Lightweight quality score ─────────────────────────────────────────────────
+
+class TestLightweightScore:
+
+    def _make_df(self, n_rows=10, n_missing=0, n_dups=0):
+        """Create a minimal test dataframe."""
+        df = pd.DataFrame({
+            "feature_0": [float(i) for i in range(n_rows)],
+            "target": [i % 2 for i in range(n_rows)],
+        })
+        if n_missing:
+            df.loc[:n_missing - 1, "feature_0"] = float("nan")
+        if n_dups:
+            df = pd.concat([df, df.iloc[:n_dups]], ignore_index=True)
+        return df
+
+    def test_clean_df_high_score(self):
+        df = self._make_df()
+        score = compute_lightweight_score(df, df.copy(), len(df),
+                                          {"feature_0": {"expected_dtype": "float64"}})
+        assert score >= 0.80
+
+    def test_many_missing_low_score(self):
+        df = self._make_df(n_missing=8)
+        score = compute_lightweight_score(df, df.copy(), 10,
+                                          {"feature_0": {"expected_dtype": "float64"}},
+                                          initial_missing=8)
+        assert score < 0.70
+
+    def test_score_in_range(self):
+        df = self._make_df(n_missing=3, n_dups=2)
+        score = compute_lightweight_score(df, df.copy(), 10,
+                                          {"feature_0": {"expected_dtype": "float64"}},
+                                          initial_missing=3)
+        assert 0.0 <= score <= 1.0

train_colab.ipynb

@@ -0,0 +1,498 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 5,
+ "metadata": {
+  "colab": {"provenance": [], "gpuType": "T4"},
+  "kernelspec": {"display_name": "Python 3", "name": "python3"},
+  "language_info": {"name": "python"},
+  "accelerator": "GPU"
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "a0",
+   "metadata": {},
+   "source": [
+    "# Data-Centric AI Agent \u2014 Training Notebook\n",
+    "**OpenEnv Hackathon 2026**\n\n",
+    "Trains a Qwen2.5-3B agent to improve ML datasets using:\n",
+    "- **Phase 1**: SFT warmup (~30 min) \u2014 teaches command grammar\n",
+    "- **Phase 2**: GRPO training (~3-5 hrs) \u2014 teaches strategy via RL\n",
+    "- **Phase 3**: Eval + reward curves\n\n",
+    "**Runtime required:** T4 GPU. `Runtime \u2192 Change runtime type \u2192 T4 GPU`"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a1",
+   "metadata": {},
+   "source": ["## Step 1 \u2014 Install dependencies (~5 min)"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\" --quiet\n",
+    "!pip install trl datasets transformers accelerate scikit-learn pandas numpy matplotlib --quiet\n",
+    "!pip install \"openenv-core[core]>=0.2.1\" --quiet\n",
+    "print('Dependencies installed')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a2",
+   "metadata": {},
+   "source": [
+    "## Step 2 \u2014 Mount Google Drive (optional but recommended)\n",
+    "Checkpoints save every 50 GRPO steps to Drive. If Colab disconnects, you can resume.\n\n",
+    "**Skip this cell if you don't want Drive integration.**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from google.colab import drive\n",
+    "drive.mount('/content/drive')\n",
+    "print('Drive mounted')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a3",
+   "metadata": {},
+   "source": ["## Step 3 \u2014 Clone / update repo + setup"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, sys, shutil\n",
+    "\n",
+    "REPO     = 'https://github.com/CelestialWorthyOfHeavenAndEarth/data-centric-env.git'\n",
+    "WORK_DIR = '/content/data-centric-env'\n",
+    "\n",
+    "# Clone or pull latest\n",
+    "if os.path.exists(f'{WORK_DIR}/pyproject.toml'):\n",
+    "    print('Repo exists, pulling latest (includes new features)...')\n",
+    "    !git -C {WORK_DIR} pull origin main\n",
+    "else:\n",
+    "    if os.path.exists(WORK_DIR):\n",
+    "        shutil.rmtree(WORK_DIR)\n",
+    "    !git clone {REPO} {WORK_DIR}\n",
+    "\n",
+    "os.chdir(WORK_DIR)\n",
+    "sys.path.insert(0, WORK_DIR)\n",
+    "print('CWD:', os.getcwd())\n",
+    "\n",
+    "# Show latest commit so we know we have the right version\n",
+    "!git log --oneline -3\n",
+    "\n",
+    "# Install as editable package (fixes relative imports)\n",
+    "!pip install -e . --quiet 2>/dev/null || echo 'pip install -e . skipped'\n",
+    "\n",
+    "# Symlink output dirs to Drive if mounted\n",
+    "DRIVE_DIR = '/content/drive/MyDrive/data-centric-training'\n",
+    "if os.path.exists('/content/drive/MyDrive'):\n",
+    "    os.makedirs(DRIVE_DIR, exist_ok=True)\n",
+    "    for d in ['data-centric-checkpoints', 'data-centric-adapter', 'logs', 'plots']:\n",
+    "        local  = os.path.join(WORK_DIR, d)\n",
+    "        remote = os.path.join(DRIVE_DIR, d)\n",
+    "        os.makedirs(remote, exist_ok=True)\n",
+    "        if not os.path.exists(local):\n",
+    "            os.symlink(remote, local)\n",
+    "            print(f'  Drive link: {d}')\n",
+    "        else:\n",
+    "            print(f'  Exists: {d}')\n",
+    "    print('Checkpoints will save to Drive')\n",
+    "else:\n",
+    "    print('Drive not mounted - checkpoints save locally only')\n",
+    "\n",
+    "print('Setup complete')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a35",
+   "metadata": {},
+   "source": [
+    "## Step 3.5 \u2014 Verify all features work\n",
+    "Runs a quick smoke test to confirm all 5 new features (query_analyst, smarter specialists,\n",
+    "dual classifier, feature importance, drift detection) are connected and working."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b35",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, sys\n",
+    "os.chdir('/content/data-centric-env')\n",
+    "sys.path.insert(0, '/content/data-centric-env')\n",
+    "\n",
+    "# Quick inline smoke test (no file needed)\n",
+    "from models import DataCentricAction\n",
+    "from server.data_centric_environment import DataCentricEnvironment\n",
+    "\n",
+    "env = DataCentricEnvironment()\n",
+    "obs = env.reset(task='task_1_easy', seed=42)\n",
+    "\n",
+    "checks = {}\n",
+    "\n",
+    "# 1. query_analyst\n",
+    "obs = env.step(DataCentricAction(message='query_analyst'))\n",
+    "checks['query_analyst'] = 'DIAGNOSIS' in obs.response and 'RECOMMENDED PLAN' in obs.response\n",
+    "\n",
+    "# 2. Smarter specialists\n",
+    "obs = env.step(DataCentricAction(message='query_cleaner'))\n",
+    "checks['smarter_specialists'] = any(k in obs.response for k in ['skew', 'Risk:', 'Reason:', '%'])\n",
+    "\n",
+    "# 3. Drift detection (apply then check)\n",
+    "obs = env.step(DataCentricAction(message='apply 1'))\n",
+    "checks['drift_detection'] = 'Distribution drift' in obs.response or 'drift' in obs.response.lower()\n",
+    "\n",
+    "# 4+5. Dual classifier + feature importance\n",
+    "obs = env.step(DataCentricAction(message='validate'))\n",
+    "checks['dual_classifier'] = 'RF Accuracy' in obs.response and 'LR Accuracy' in obs.response\n",
+    "checks['feature_importance'] = 'Feature importance' in obs.response\n",
+    "\n",
+    "print('\\n=== Feature Smoke Test ===')\n",
+    "all_pass = True\n",
+    "for name, passed in checks.items():\n",
+    "    status = 'PASS' if passed else 'FAIL'\n",
+    "    if not passed: all_pass = False\n",
+    "    print(f'  {status}: {name}')\n",
+    "\n",
+    "if all_pass:\n",
+    "    print('\\nAll 5 features verified OK - ready to train!')\n",
+    "else:\n",
+    "    print('\\nWARNING: Some features failed - check git pull ran correctly')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a4",
+   "metadata": {},
+   "source": ["## Step 4 \u2014 Start environment server"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import subprocess, time, requests, os, sys\n",
+    "os.chdir('/content/data-centric-env')\n",
+    "sys.path.insert(0, '/content/data-centric-env')\n",
+    "\n",
+    "# Kill any existing server on port 8000\n",
+    "!fuser -k 8000/tcp 2>/dev/null || true\n",
+    "time.sleep(2)\n",
+    "\n",
+    "server_proc = subprocess.Popen(\n",
+    "    ['python', '-m', 'uvicorn', 'server.app:app',\n",
+    "     '--host', '0.0.0.0', '--port', '8000', '--log-level', 'warning'],\n",
+    "    stdout=subprocess.PIPE, stderr=subprocess.PIPE\n",
+    ")\n",
+    "\n",
+    "for i in range(30):\n",
+    "    try:\n",
+    "        if requests.get('http://localhost:8000/health', timeout=2).status_code == 200:\n",
+    "            print(f'Server ready after {i+1}s')\n",
+    "            break\n",
+    "    except:\n",
+    "        time.sleep(1)\n",
+    "else:\n",
+    "    server_proc.terminate()\n",
+    "    out, err = server_proc.communicate()\n",
+    "    print('STDOUT:', out.decode()[-1000:])\n",
+    "    print('STDERR:', err.decode()[-1000:])\n",
+    "    raise RuntimeError('Server failed to start in 30s - see output above')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a5",
+   "metadata": {},
+   "source": ["## Step 5 \u2014 Generate SFT warmup data"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.chdir('/content/data-centric-env')\n",
+    "\n",
+    "# Always regenerate if we updated the code (new query_analyst strategies)\n",
+    "# Skip only if the file is recent (within last 10 minutes)\n",
+    "import time\n",
+    "sft_path = 'sft_data.jsonl'\n",
+    "regen = True\n",
+    "if os.path.exists(sft_path):\n",
+    "    age_minutes = (time.time() - os.path.getmtime(sft_path)) / 60\n",
+    "    count = sum(1 for _ in open(sft_path))\n",
+    "    print(f'sft_data.jsonl exists: {count} examples, {age_minutes:.0f} min old')\n",
+    "    # Regenerate if old file missing query_analyst strategies (<= 18 strategies)\n",
+    "    regen = count < 20000  # new version has ~21 strategies\n",
+    "\n",
+    "if regen:\n",
+    "    print('Generating SFT data (includes query_analyst strategies)...')\n",
+    "    !python sft_generator.py\n",
+    "    print('SFT data generated')\n",
+    "else:\n",
+    "    print('SFT data is up to date - skipping')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a6",
+   "metadata": {},
+   "source": ["## Step 6 \u2014 Load model (Qwen2.5-3B-Instruct, 4-bit)"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, sys\n",
+    "os.chdir('/content/data-centric-env')\n",
+    "sys.path.insert(0, '/content/data-centric-env')\n",
+    "\n",
+    "from unsloth import FastLanguageModel\n",
+    "\n",
+    "MODEL_NAME = 'Qwen/Qwen2.5-3B-Instruct'\n",
+    "MAX_SEQ_LENGTH = 1024\n",
+    "\n",
+    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
+    "    model_name=MODEL_NAME,\n",
+    "    max_seq_length=MAX_SEQ_LENGTH,\n",
+    "    load_in_4bit=True,\n",
+    "    dtype=None,\n",
+    ")\n",
+    "\n",
+    "model = FastLanguageModel.get_peft_model(\n",
+    "    model,\n",
+    "    r=16,\n",
+    "    target_modules=['q_proj', 'v_proj', 'k_proj', 'o_proj'],\n",
+    "    lora_alpha=32,\n",
+    "    lora_dropout=0,\n",
+    "    bias='none',\n",
+    "    use_gradient_checkpointing='unsloth',\n",
+    "    random_state=42,\n",
+    ")\n",
+    "\n",
+    "import torch\n",
+    "print(f'Model loaded: {MODEL_NAME}')\n",
+    "print(f'VRAM: {torch.cuda.memory_allocated()/1e9:.1f} GB')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a7",
+   "metadata": {},
+   "source": ["## Step 7 \u2014 Phase 1: SFT Warmup (~30 min)"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, sys\n",
+    "os.chdir('/content/data-centric-env')\n",
+    "sys.path.insert(0, '/content/data-centric-env')\n",
+    "\n",
+    "from train_data_centric import run_sft_warmup\n",
+    "model = run_sft_warmup(model, tokenizer)\n",
+    "print('SFT warmup complete')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a8",
+   "metadata": {},
+   "source": [
+    "## Step 8 \u2014 Phase 2: GRPO Training (~3-5 hrs on T4)\n\n",
+    "Agent learns **when** to use each command via reinforcement learning.\n",
+    "Checkpoints save every 50 steps. If Colab disconnects, re-run all cells \u2014\n",
+    "it auto-resumes from the latest checkpoint.\n\n",
+    "**What the agent will learn:**\n",
+    "- Start with `query_analyst` to get a prioritised plan\n",
+    "- Use the Agreement signal to detect overfitting\n",
+    "- Focus cleaning on high-importance features\n",
+    "- Undo bad applies when drift is HIGH"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, sys, glob\n",
+    "os.chdir('/content/data-centric-env')\n",
+    "sys.path.insert(0, '/content/data-centric-env')\n",
+    "\n",
+    "from train_data_centric import run_grpo_training\n",
+    "\n",
+    "# Auto-detect checkpoint to resume from\n",
+    "ckpt_dir = './data-centric-checkpoints'\n",
+    "resume_from = None\n",
+    "if os.path.exists(ckpt_dir):\n",
+    "    checkpoints = sorted(glob.glob(f'{ckpt_dir}/checkpoint-*'))\n",
+    "    if checkpoints:\n",
+    "        resume_from = checkpoints[-1]\n",
+    "        print(f'Resuming from: {resume_from}')\n",
+    "    else:\n",
+    "        print('No checkpoint found - starting fresh')\n",
+    "\n",
+    "model = run_grpo_training(model, tokenizer, resume_from_checkpoint=resume_from)\n",
+    "print('GRPO training complete')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a9",
+   "metadata": {},
+   "source": ["## Step 9 \u2014 Save trained model"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, sys\n",
+    "os.chdir('/content/data-centric-env')\n",
+    "sys.path.insert(0, '/content/data-centric-env')\n",
+    "\n",
+    "from train_data_centric import save_model\n",
+    "save_model(model, tokenizer)\n",
+    "print('Model saved to ./data-centric-adapter/')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a10",
+   "metadata": {},
+   "source": ["## Step 10 \u2014 Plot reward curves"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.chdir('/content/data-centric-env')\n",
+    "\n",
+    "if os.path.exists('logs/training.jsonl'):\n",
+    "    lines = sum(1 for _ in open('logs/training.jsonl'))\n",
+    "    print(f'Training log: {lines} episodes recorded')\n",
+    "    !python plot_rewards.py --log logs/training.jsonl --out plots/\n",
+    "    from IPython.display import Image, display\n",
+    "    for f in ['reward_curve.png', 'success_rate.png', 'accuracy_gain.png', 'curriculum.png']:\n",
+    "        path = f'plots/{f}'\n",
+    "        if os.path.exists(path):\n",
+    "            print(f'\\n--- {f} ---')\n",
+    "            display(Image(path))\n",
+    "else:\n",
+    "    print('No training log found yet - run Step 8 first')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a11",
+   "metadata": {},
+   "source": ["## Step 11 \u2014 Evaluate trained agent vs random vs heuristic"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, json\n",
+    "os.chdir('/content/data-centric-env')\n",
+    "\n",
+    "!python eval_data_centric.py\n",
+    "\n",
+    "if os.path.exists('eval_results.json'):\n",
+    "    with open('eval_results.json') as f:\n",
+    "        results = json.load(f)\n",
+    "    print('\\n=== Eval Results ===')\n",
+    "    for k, v in results.items():\n",
+    "        print(f'  {k}: {v}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a12",
+   "metadata": {},
+   "source": [
+    "## Step 12 \u2014 Push results to GitHub\n\n",
+    "Run this **only after training + eval are done**."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.chdir('/content/data-centric-env')\n",
+    "\n",
+    "from getpass import getpass\n",
+    "token = getpass('GitHub token (repo write access): ')\n",
+    "repo_url = f'https://{token}@github.com/CelestialWorthyOfHeavenAndEarth/data-centric-env.git'\n",
+    "\n",
+    "!git config user.email 'training@colab.run'\n",
+    "!git config user.name 'Colab Training'\n",
+    "!git remote set-url origin {repo_url}\n",
+    "!git add plots/ logs/ eval_results.json 2>/dev/null || true\n",
+    "!git commit -m 'Add training results: reward curves + eval'\n",
+    "!git push origin main\n",
+    "print('Results pushed to GitHub')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a13",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## Done!\n\n",
+    "| Output | Location |\n",
+    "|--------|----------|\n",
+    "| Trained adapter | `./data-centric-adapter/` |\n",
+    "| Training log | `logs/training.jsonl` |\n",
+    "| Reward curves | `plots/*.png` |\n",
+    "| Eval results | `eval_results.json` |"
+   ]
+  }
+ ]
+}

train_data_centric.py

@@ -0,0 +1,689 @@
+# pip install trl unsloth transformers torch requests
+# pip install matplotlib openenv-core scikit-learn pandas numpy datasets
+
+import os
+import json
+import random
+import time
+import signal
+import subprocess
+import requests
+import torch
+from collections import deque
+from statistics import mean
+from datasets import Dataset
+from unsloth import FastLanguageModel
+from trl import SFTTrainer, SFTConfig, GRPOConfig, GRPOTrainer
+
+# WebSocket client for stateful episode rollouts
+sys_path_root = os.path.dirname(os.path.abspath(__file__))
+import sys
+if sys_path_root not in sys.path:
+    sys.path.insert(0, sys_path_root)
+from client import DataCentricEnv
+from models import DataCentricAction
+
+# ════════════════════════════════════════════════════════
+# CONSTANTS
+# ════════════════════════════════════════════════════════
+
+# ENV_URL: set this to your HF Space URL when running as an HF Job
+# e.g. export ENV_URL="https://aswini-kumar-data-centric-env.hf.space"
+BASE_URL = os.environ.get("ENV_URL", "http://localhost:8000")
+MODEL_NAME = "Qwen/Qwen2.5-3B-Instruct"
+MAX_SEQ_LENGTH = 1024
+LOAD_IN_4BIT = True
+
+VALID_COMMANDS = [
+    "inspect_dataset", "inspect_model", "query_analyst",
+    "query_cleaner", "query_augmenter", "query_balancer", "query_validator",
+    "apply", "reject", "undo", "validate", "submit",
+]
+
+SYSTEM_PROMPT = """You are a Data-Centric AI agent. Your job is to improve a \
+Machine learning dataset so a fixed classifier achieves higher accuracy.
+
+STRATEGY — use this order:
+1. query_analyst to get a prioritised action plan (costs 1 budget, worth it)
+2. inspect_dataset to understand the data
+3. query the recommended specialist (query_cleaner, query_augmenter, query_balancer)
+4. apply the best recommendation by number (e.g. apply 1)
+5. validate to check if accuracy improved
+6. repeat until you hit the target or run low on budget
+7. submit to finalize
+
+IMPORTANT RULES:
+- Start with query_analyst — it tells you the biggest problem to fix first.
+- Always query a specialist before applying. Never apply without querying first.
+- Check the Agreement signal after validate: DISAGREE means possible overfitting.
+- Validate after every 2-3 applies to track progress.
+- Do not spam validate — it costs budget after 3 uses.
+- query_validator costs 2 budget — use only when suspicious of data quality.
+- submit when accuracy >= target or budget < 5.
+
+Reply with exactly ONE command per message. No explanation. Just the command."""
+
+
+def build_user_prompt(obs: dict) -> str:
+    improvement_needed = obs.get("target_accuracy", 0) - obs.get("current_accuracy", 0)
+    return (
+        f"Current situation:\n"
+        f"Accuracy: {obs.get('current_accuracy', 0):.1%} → "
+        f"Target: {obs.get('target_accuracy', 0):.1%}\n"
+        f"Still need: {max(0, improvement_needed):.1%} improvement\n"
+        f"Quality score: {obs.get('estimated_quality', 0):.2f}/1.00\n"
+        f"Rows preserved: {obs.get('rows_preserved_pct', 1.0):.1%}\n"
+        f"Budget remaining: {obs.get('budget_remaining', 0)} steps\n"
+        f"Free validates left: {obs.get('validate_calls_remaining', 0)}\n"
+        f"Active query session: {obs.get('active_session', 'none')}\n\n"
+        f"Last result:\n{str(obs.get('response', ''))[:400]}\n\n"
+        f"What is your next action? (one command only)"
+    )
+
+
+# ════════════════════════════════════════════════════════
+# SERVER MANAGEMENT
+# ════════════════════════════════════════════════════════
+
+def start_server() -> subprocess.Popen:
+    """Start the environment server as a subprocess."""
+    proc = subprocess.Popen(
+        ["uvicorn", "server.app:app", "--host", "0.0.0.0", "--port", "8000"],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+    # Poll until ready (max 30 seconds)
+    for i in range(30):
+        try:
+            r = requests.get(f"{BASE_URL}/health", timeout=1)
+            if r.status_code == 200:
+                print(f"Server ready after {i + 1}s")
+                return proc
+        except Exception:
+            pass
+        time.sleep(1)
+    proc.terminate()
+    raise RuntimeError("Environment server failed to start in 30 seconds")
+
+
+def stop_server(proc: subprocess.Popen):
+    proc.terminate()  # cross-platform (SIGTERM on Linux, TerminateProcess on Windows)
+    proc.wait()
+    print("Server stopped.")
+
+
+# ════════════════════════════════════════════════════════
+# MODEL SETUP
+# ════════��═══════════════════════════════════════════════
+
+def load_model():
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name=MODEL_NAME,
+        max_seq_length=MAX_SEQ_LENGTH,
+        load_in_4bit=LOAD_IN_4BIT,
+        dtype=None,
+    )
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r=16,
+        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
+        lora_alpha=32,
+        lora_dropout=0,
+        bias="none",
+        use_gradient_checkpointing="unsloth",
+        random_state=42,
+    )
+    return model, tokenizer
+
+
+# ════════════════════════════════════════════════════════
+# PHASE 1 — SFT WARMUP
+# ════════════════════════════════════════════════════════
+
+def run_sft_warmup(model, tokenizer):
+    """
+    1 epoch of SFT on heuristic trajectories.
+    Teaches model valid command format before GRPO starts.
+    Without this, model outputs gibberish and gets zero reward.
+    """
+    print("\n=== PHASE 1: SFT WARMUP ===")
+    
+    if os.path.exists("./sft-checkpoint"):
+        model.load_adapter("./sft-checkpoint")
+        print("Loaded existing SFT checkpoint — skipping warmup.")
+        return model
+
+    if not os.path.exists("sft_data.jsonl"):
+        print("sft_data.jsonl not found. Run sft_generator.py first.")
+        raise FileNotFoundError("sft_data.jsonl missing")
+
+    raw = [json.loads(l) for l in open("sft_data.jsonl", encoding="utf-8")]
+    print(f"Loaded {len(raw)} SFT examples")
+
+    def format_example(ex):
+        messages = [
+            {"role": "system",    "content": SYSTEM_PROMPT},
+            {"role": "user",      "content": ex["prompt"]},
+            {"role": "assistant", "content": ex["response"]},
+        ]
+        return {
+            "text": tokenizer.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=False
+            )
+        }
+
+    sft_dataset = Dataset.from_list([format_example(ex) for ex in raw])
+
+    sft_trainer = SFTTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=sft_dataset,
+        args=SFTConfig(
+            output_dir="./sft-checkpoint",
+            num_train_epochs=1,
+            per_device_train_batch_size=4,
+            gradient_accumulation_steps=4,
+            learning_rate=2e-5,
+            warmup_steps=10,
+            logging_steps=5,
+            save_strategy="no",
+            report_to="none",
+            max_seq_length=MAX_SEQ_LENGTH,
+        ),
+    )
+    sft_trainer.train()
+    print("SFT warmup complete.\n")
+    return model
+
+
+# ════════════════════════════════════════════════════════
+# CURRICULUM SCHEDULER
+# ════════════════════════════════════════════════════════
+
+class CurriculumScheduler:
+    """
+    Advances curriculum level when the agent reliably solves the current task.
+
+    Advancement criterion: >= threshold success rate over a rolling window of episodes.
+    Uses a smoothed window to avoid premature advancement on lucky streaks.
+
+    Levels: 0=tutorial, 1=easy, 2=medium, 3=hard
+
+    Design rationale:
+    - Step-count based scheduling causes premature advancement (catastrophic forgetting)
+      or stalling (wasted compute) because it ignores actual agent performance.
+    - Success-rate based scheduling ensures the agent genuinely masters a level
+      before seeing harder tasks, matching curriculum RL best practices.
+    - Window resets after each advancement so the agent must prove itself again.
+    """
+
+    TASKS = ["task_0_tutorial", "task_1_easy", "task_2_medium", "task_3_hard"]
+    LEVEL_LABELS = ["tutorial", "easy", "medium", "hard"]
+
+    def __init__(self, window: int = 30, threshold: float = 0.60):
+        """
+        Args:
+            window: Number of episodes to evaluate before considering advancement.
+                    30 gives a stable estimate without waiting too long.
+            threshold: Required success rate (0.0–1.0) to advance to next level.
+                       0.60 = agent must solve 60% of episodes to advance.
+                       Lower than classic 0.75 because 3B models learn slower.
+        """
+        self.current_level = 0
+        self.window = window
+        self.threshold = threshold
+        self.recent_successes: deque = deque(maxlen=window)
+        self.global_step = 0          # total episodes recorded (for logging)
+        self.level_history: list = [] # (episode, level) pairs for plotting
+
+    def record_episode(self, reached_target: bool, accuracy_gain: float = 0.0):
+        """Call after every episode completes."""
+        self.recent_successes.append(float(reached_target))
+        self.global_step += 1
+        if self.should_advance():
+            self.advance()
+
+    def get_task(self) -> str:
+        """Return the current training task name."""
+        return self.TASKS[self.current_level]
+
+    def current_success_rate(self) -> float:
+        if not self.recent_successes:
+            return 0.0
+        return sum(self.recent_successes) / len(self.recent_successes)
+
+    def should_advance(self) -> bool:
+        """Only advance if we have enough data and consistently exceed threshold."""
+        return (
+            len(self.recent_successes) >= self.window
+            and self.current_success_rate() >= self.threshold
+            and self.current_level < len(self.TASKS) - 1
+        )
+
+    def advance(self):
+        if self.current_level < len(self.TASKS) - 1:
+            print(
+                f"\n[Curriculum] ▶ Level {self.current_level} ({self.TASKS[self.current_level]}) "
+                f"→ Level {self.current_level + 1} ({self.TASKS[self.current_level + 1]})\n"
+                f"  Success rate over last {self.window} episodes: "
+                f"{self.current_success_rate():.1%} (threshold: {self.threshold:.0%})\n"
+                f"  Total episodes: {self.global_step}"
+            )
+            self.level_history.append((self.global_step, self.current_level))
+            self.current_level += 1
+            self.recent_successes.clear()  # reset window after advancing
+
+    def stage_label(self) -> str:
+        return self.LEVEL_LABELS[self.current_level]
+
+    # Backward-compat: record_improvement still works for old callers
+    def record_improvement(self, improvement: float):
+        self.record_episode(reached_target=(improvement > 0.05))
+        self.global_step = self.global_step  # already incremented above
+
+
+
+
+# ════════════════════════════════════════════════════════
+# REWARD COMPUTATION
+# ════════════════════════════════════════════════════════
+
+def compute_rewards(
+    obs_before: dict,
+    obs_after: dict,
+    response_text: str,
+    action_history: list,
+) -> dict:
+    """
+    4 completely independent reward components.
+    Each measures a different aspect of agent behavior.
+    No single component can be maximized without solving the real task.
+    """
+    # Component 1: Environment accuracy reward (from environment)
+    env_reward = obs_after.get("reward", 0.0)
+
+    # Component 2: Format reward — did model output a valid command?
+    is_valid = any(
+        response_text.strip().startswith(cmd) for cmd in VALID_COMMANDS
+    )
+    format_reward = 0.10 if is_valid else -0.10
+
+    # Component 3: Strategy reward — did model follow smart workflow?
+    strategy_reward = 0.0
+    if len(action_history) >= 1:
+        prev = action_history[-1]
+        curr = response_text.strip()
+
+        if prev.startswith("query_") and curr.startswith("apply"):
+            strategy_reward = 0.05   # query then apply — correct pattern
+        elif curr.startswith("apply") and not any(
+            a.startswith("query_") for a in action_history[-3:]
+        ):
+            strategy_reward = -0.04  # apply without recent query — bad pattern
+        elif prev.startswith("apply") and curr == "validate":
+            strategy_reward = 0.03   # validate after apply — good pattern
+        elif curr == "submit" and not any(
+            a == "validate" for a in action_history
+        ):
+            strategy_reward = -0.10  # submit without ever validating — bad
+
+    # Component 4: Preservation reward — did model lose good data?
+    rows_pct = obs_after.get("rows_preserved_pct", 1.0)
+    if rows_pct >= 0.90:   preservation_reward = 0.05
+    elif rows_pct >= 0.80: preservation_reward = 0.02
+    elif rows_pct >= 0.70: preservation_reward = 0.00
+    elif rows_pct >= 0.50: preservation_reward = -0.10
+    else:                  preservation_reward = -0.40
+
+    total = env_reward + format_reward + strategy_reward + preservation_reward
+
+    return {
+        "total":        total,
+        "env":          env_reward,
+        "format":       format_reward,
+        "strategy":     strategy_reward,
+        "preservation": preservation_reward,
+    }
+
+
+# ════════════════════════════════════════════════════════
+# EPISODE ROLLOUT
+# ════════════════════════════════════════════════════════
+
+def obs_to_dict(obs_obj) -> dict:
+    """Convert DataCentricObservation to dict for compatibility with reward logic."""
+    if isinstance(obs_obj, dict):
+        return obs_obj
+    return {
+        "response":                obs_obj.response,
+        "current_accuracy":        obs_obj.current_accuracy,
+        "baseline_accuracy":       obs_obj.baseline_accuracy,
+        "target_accuracy":         obs_obj.target_accuracy,
+        "estimated_quality":       obs_obj.estimated_quality,
+        "dataset_shape":           obs_obj.dataset_shape,
+        "rows_preserved_pct":      obs_obj.rows_preserved_pct,
+        "budget_remaining":        obs_obj.budget_remaining,
+        "step_number":             obs_obj.step_number,
+        "max_steps":               obs_obj.max_steps,
+        "active_session":          obs_obj.active_session,
+        "validate_calls_remaining":obs_obj.validate_calls_remaining,
+        "done":                    obs_obj.done,
+        "reward":                  obs_obj.reward,
+    }
+
+
+def run_episode(
+    model, tokenizer, task: str, seed: int
+) -> tuple:
+    """
+    Run one complete episode using the WebSocket client (stateful session).
+    Each reset+step sequence maintains the same env instance on the server.
+    Returns: (prompts, responses, rewards) for GRPO training.
+    """
+    prompts, responses, rewards = [], [], []
+    action_history = []
+
+    try:
+        with DataCentricEnv(base_url=BASE_URL).sync() as env:
+            reset_result = env.reset(task=task, seed=seed)
+            obs = obs_to_dict(reset_result.observation)
+
+            while not obs.get("done", False):
+                # Build chat prompt
+                messages = [
+                    {"role": "system", "content": SYSTEM_PROMPT},
+                    {"role": "user",   "content": build_user_prompt(obs)},
+                ]
+                input_ids = tokenizer.apply_chat_template(
+                    messages,
+                    return_tensors="pt",
+                    max_length=MAX_SEQ_LENGTH - 60,
+                    truncation=True,
+                    add_generation_prompt=True,
+                ).to(model.device)
+
+                # Generate — commands are short, 50 tokens max
+                with torch.no_grad():
+                    output_ids = model.generate(
+                        input_ids,
+                        max_new_tokens=50,
+                        temperature=0.8,
+                        do_sample=True,
+                        pad_token_id=tokenizer.eos_token_id,
+                    )
+
+                response_text = tokenizer.decode(
+                    output_ids[0][input_ids.shape[1]:],
+                    skip_special_tokens=True,
+                ).strip().split("\n")[0].strip()[:200]
+
+                obs_before = obs
+                try:
+                    step_result = env.step(DataCentricAction(message=response_text))
+                    obs = obs_to_dict(step_result.observation)
+                except Exception as e:
+                    obs = {**obs, "done": True, "reward": -0.05,
+                           "response": f"Step error: {e}"}
+
+                reward_dict = compute_rewards(
+                    obs_before, obs, response_text, action_history
+                )
+                prompts.append(build_user_prompt(obs_before))
+                responses.append(response_text)
+                rewards.append(reward_dict)
+                action_history.append(response_text)
+
+    except Exception as e:
+        print(f"Episode error (task={task}, seed={seed}): {e}")
+        return [], [], []
+
+    return prompts, responses, rewards
+
+
+# ════════════════════════════════════════════════════════
+# LOGGING
+# ════════════════════════════════════════════════════════
+
+training_log = []
+
+
+def log_training_step(
+    step: int, all_episodes: list, scheduler: CurriculumScheduler
+):
+    """Log metrics and sample generations every 10 steps."""
+    all_final_rewards = []
+    all_reward_components: dict = {
+        "env": [], "format": [], "strategy": [], "preservation": []
+    }
+    format_hits = 0
+    strategy_hits = 0
+    total_actions = 0
+
+    for prompts, responses, rewards in all_episodes:
+        if not rewards:
+            continue
+        all_final_rewards.append(rewards[-1]["total"])
+        for r in rewards:
+            for k in all_reward_components:
+                all_reward_components[k].append(r[k])
+            if r["format"] > 0:   format_hits += 1
+            if r["strategy"] > 0: strategy_hits += 1
+            total_actions += 1
+
+    if not all_final_rewards:
+        return
+
+    entry = {
+        "step":                  step,
+        "stage":                 scheduler.stage_label(),
+        "task":                  scheduler.get_task(),
+        "mean_total_reward":     mean(all_final_rewards),
+        "mean_env_reward":       mean(all_reward_components["env"]),
+        "mean_format_reward":    mean(all_reward_components["format"]),
+        "mean_strategy_reward":  mean(all_reward_components["strategy"]),
+        "mean_preservation_reward": mean(all_reward_components["preservation"]),
+        "format_rate":           format_hits / max(total_actions, 1),
+        "strategy_rate":         strategy_hits / max(total_actions, 1),
+    }
+    training_log.append(entry)
+    json.dump(training_log, open("training_log.json", "w"), indent=2)
+
+    print(
+        f"Step {step:4d} | Stage: {entry['stage']:8s} | "
+        f"Reward: {entry['mean_total_reward']:+.3f} | "
+        f"Format: {entry['format_rate']:.0%} | "
+        f"Strategy: {entry['strategy_rate']:.0%}"
+    )
+
+    # Sample 3 generations for inspection
+    if step % 10 == 0:
+        samples = []
+        for p_ep, r_ep, rw_ep in all_episodes[:3]:
+            if p_ep and r_ep:
+                samples.append({
+                    "step":          step,
+                    "response":      r_ep[-1],
+                    "reward":        rw_ep[-1]["total"],
+                    "reward_detail": rw_ep[-1],
+                })
+        with open("generations.jsonl", "a", encoding="utf-8") as f:
+            for s in samples:
+                f.write(json.dumps(s) + "\n")
+
+
+def log_episode_jsonl(
+    episode: int, task: str, level: int, reward: float,
+    accuracy_gain: float, steps_used: int, success: bool,
+    log_path: str = "logs/training.jsonl",
+):
+    """Write one episode record to JSONL log (read by plot_rewards.py)."""
+    import os as _os
+    _os.makedirs(_os.path.dirname(log_path), exist_ok=True)
+    entry = {
+        "ts":            time.time(),
+        "episode":       episode,
+        "task":          task,
+        "level":         level,
+        "reward":        round(reward, 4),
+        "accuracy_gain": round(accuracy_gain, 4),
+        "steps_used":    steps_used,
+        "success":       success,
+    }
+    with open(log_path, "a", encoding="utf-8") as f:
+        f.write(json.dumps(entry) + "\n")
+
+
+# ════════════════════════════════════════════════════════
+# GRPO TRAINING LOOP
+# ════════════════════════════════════════════════════════
+
+def run_grpo_training(model, tokenizer, resume_from_checkpoint=None):
+    print("\n=== PHASE 2: GRPO TRAINING ===")
+    if resume_from_checkpoint:
+        print(f"Resuming from checkpoint: {resume_from_checkpoint}")
+
+    scheduler = CurriculumScheduler()
+
+    grpo_config = GRPOConfig(
+        output_dir="./data-centric-checkpoints",
+        num_train_epochs=3,
+        per_device_train_batch_size=4,
+        gradient_accumulation_steps=4,
+        learning_rate=5e-6,
+        warmup_steps=20,
+        logging_steps=10,
+        save_steps=50,
+        num_generations=4,
+        max_completion_length=50,   # renamed from max_new_tokens in TRL ≥0.15
+        max_prompt_length=900,
+        report_to="none",
+    )
+
+    def reward_fn(completions, prompts=None, **kwargs):
+        """
+        Reward function called by GRPOTrainer.
+        Runs live episodes and returns total reward for each completion.
+        """
+        batch_rewards = []
+        episodes_this_batch = []
+
+        for completion in completions:
+            # Capture task BEFORE running episode so log reflects what was run
+            task = scheduler.get_task()
+            seed = random.randint(0, 9999)
+
+            prompts_ep, responses_ep, rewards_ep = run_episode(
+                model, tokenizer, task, seed
+            )
+
+            if rewards_ep:
+                final_reward = sum(r["total"] for r in rewards_ep)
+                accuracy_gain = sum(r["env"] for r in rewards_ep)
+                success = accuracy_gain > 0.05
+                # Update curriculum using success-rate based scheduler
+                scheduler.record_episode(
+                    reached_target=success,
+                    accuracy_gain=accuracy_gain,
+                )
+            else:
+                final_reward = -0.10
+                accuracy_gain = 0.0
+                success = False
+                scheduler.record_episode(reached_target=False, accuracy_gain=0.0)
+
+            # Write per-episode JSONL record for plot_rewards.py
+            log_episode_jsonl(
+                episode=scheduler.global_step,
+                task=task,
+                level=scheduler.current_level,
+                reward=final_reward,
+                accuracy_gain=accuracy_gain,
+                steps_used=len(rewards_ep) if rewards_ep else 0,
+                success=success,
+            )
+
+            batch_rewards.append(final_reward)
+            episodes_this_batch.append((prompts_ep, responses_ep, rewards_ep))
+
+        # Log every 10 calls
+        if scheduler.global_step % 10 == 0:
+            log_training_step(
+                scheduler.global_step,
+                episodes_this_batch,
+                scheduler,
+            )
+
+        return batch_rewards
+
+    # GRPOTrainer needs a prompt dataset — use SFT data as source
+    raw = [json.loads(l) for l in open("sft_data.jsonl", encoding="utf-8")]
+    grpo_dataset = Dataset.from_list([
+        {"prompt": ex["prompt"]} for ex in raw
+    ])
+
+    trainer = GRPOTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        reward_funcs=[reward_fn],
+        args=grpo_config,
+        train_dataset=grpo_dataset,
+    )
+    trainer.train(resume_from_checkpoint=resume_from_checkpoint)
+    print("GRPO training complete.\n")
+    return model
+
+
+# ════════════════════════════════════════════════════════
+# SAVE MODEL
+# ════════════════════════════════════════════════════════
+
+def save_model(model, tokenizer):
+    print("Saving model...")
+
+    # Save LoRA adapter (safe for 4-bit, fast)
+    model.save_pretrained("./data-centric-adapter")
+    tokenizer.save_pretrained("./data-centric-adapter")
+    print("Adapter saved to ./data-centric-adapter")
+
+    # Save merged 16-bit for inference
+    # IMPORTANT: use unsloth's method — NOT naive merge_and_unload()
+    # Naive merge on 4-bit model corrupts weights
+    model.save_pretrained_merged(
+        "./data-centric-merged",
+        tokenizer,
+        save_method="merged_16bit",
+    )
+    print("Merged model saved to ./data-centric-merged")
+    print("Test inference immediately before demo.")
+
+
+# ════════════════════════════════════════════════════════
+# MAIN
+# ════════════════════════════════════════════════════════
+
+if __name__ == "__main__":
+    # Ensure SFT warmup data exists
+    if not os.path.exists("sft_data.jsonl"):
+        print("Generating SFT data first...")
+        subprocess.run(["python", "sft_generator.py"], check=True)
+
+    # Start environment server
+    server_proc = start_server()
+
+    try:
+        # Load base model with LoRA
+        model, tokenizer = load_model()
+
+        # Phase 1: SFT warmup — teaches valid command grammar
+        model = run_sft_warmup(model, tokenizer)
+
+        # Phase 2: GRPO — improves strategy via environment reward
+        model = run_grpo_training(model, tokenizer)
+
+        # Save adapter + merged 16-bit
+        save_model(model, tokenizer)
+
+        print("\nTraining complete. Run eval_data_centric.py next.")
+
+    finally:
+        stop_server(server_proc)