README.md

---

title: Data Cleaning Environment
emoji: 🧹
colorFrom: blue
colorTo: purple
sdk: docker
app_port: 7860
base_path: /web
---

<div align="center">

# 🧹 Data Cleaning Environment

### A Reinforcement Learning Benchmark for Autonomous Data Cleaning Agents

[![Python](https://img.shields.io/badge/Python-3.12+-3776AB?style=for-the-badge&logo=python&logoColor=white)](https://www.python.org/)
[![OpenEnv](https://img.shields.io/badge/OpenEnv-Compatible-FF6B35?style=for-the-badge)](https://github.com/meta-pytorch/OpenEnv)
[![Pydantic](https://img.shields.io/badge/Pydantic-v2-E92063?style=for-the-badge&logo=pydantic&logoColor=white)](https://docs.pydantic.dev/)
[![FastAPI](https://img.shields.io/badge/FastAPI-WebSocket-009688?style=for-the-badge&logo=fastapi&logoColor=white)](https://fastapi.tiangolo.com/)
[![Docker](https://img.shields.io/badge/Docker-Ready-2496ED?style=for-the-badge&logo=docker&logoColor=white)](https://www.docker.com/)
[![HuggingFace](https://img.shields.io/badge/HuggingFace-Deployable-FFD21E?style=for-the-badge&logo=huggingface&logoColor=black)](https://huggingface.co/)
[![License](https://img.shields.io/badge/License-MIT-green?style=for-the-badge)](LICENSE)

<br/>

> **An OpenEnv-compatible reinforcement learning environment where an LLM agent receives a dirty CSV dataset and must autonomously fix type errors, outliers, missing values, and schema inconsistencies to match a hidden ground truth — step by step.**

<br/>

```

┌──────────────────────────────────────────────────────────────────┐

│   Dirty CSV  →  Agent Observes  →  Issues CleanAction  →  Reward │

│                                                                  │

│   "N/A"  →  FILL_MISSING(median)  →  Score ↑  →  +0.12 reward  │

│   "2099" →  SET_VALUE(row=3,"2024-01-15")  →  Score ↑  →  +0.08 │

│   "  bob" → STANDARDIZE_COL("name")  →  Score ↑  →  +0.05       │

└──────────────────────────────────────────────────────────────────┘

```

</div>

---

## 📑 Table of Contents

- [Overview](#-overview)
- [Architecture](#-architecture)
- [Project Structure](#-project-structure)
- [Tasks](#-tasks)
- [Action Space](#-action-space)
- [Observation Space](#-observation-space)
- [Reward Function](#-reward-function)
- [Quick Start](#-quick-start)
- [Running Inference](#-running-inference)
- [Environment API](#-environment-api)
- [Configuration](#-configuration)
- [Deployment](#-deployment)
- [Development & Testing](#-development--testing)
- [Troubleshooting](#-troubleshooting)

---

## 🌟 Overview

The **Data Cleaning Environment** is a structured RL benchmark where an LLM-powered agent must clean tabular datasets. The environment wraps a FastAPI WebSocket server following the [OpenEnv](https://github.com/meta-pytorch/OpenEnv) protocol, making it compatible with any OpenEnv-based training or evaluation framework.

### Why This Matters

Real-world data pipelines spend 60–80% of their time on data cleaning. This environment trains agents to:

- **Detect** type errors, outliers, missing values, and schema inconsistencies
- **Reason** about which fix is most impactful at each step
- **Self-correct** from informative error feedback
- **Terminate** efficiently without over-cleaning

### Key Properties

| Property | Value |
|---|---|
| Protocol | OpenEnv (WebSocket + HTTP) |
| Action Space | Discrete (5 command types) |
| Observation | Full CSV state + grader feedback |
| Episode Structure | Reset → N × Step → Done |
| Concurrency | ✅ Multiple simultaneous sessions |
| State Management | Server-side, fully isolated per session |

---

## 🏗️ Architecture

```

┌─────────────────────────────────────────────────────────────────────┐

│                         Agent (LLM / RL Policy)                     │

│                  Qwen2.5-72B / Mistral / Custom Model               │

└────────────────────────┬───────────────────────────────┬────────────┘

                         │ CleanAction (JSON)             │ CleanObservation

                         ▼                               │

┌────────────────────────────────────────────────────────┴────────────┐

│                      DataCleaningEnv (client.py)                     │

│               OpenEnv EnvClient[CleanAction, CleanObservation, dict] │

│                   WebSocket persistent connection                     │

└────────────────────────┬────────────────────────────────────────────┘

                         │  WebSocket /ws

                         ▼

┌─────────────────────────────────────────────────────────────────────┐

│                      FastAPI Server (server/app.py)                  │

│                  HTTP + WebSocket endpoints, sessions                │

└────────────────────────┬────────────────────────────────────────────┘

                         │

                         ▼

┌─────────────────────────────────────────────────────────────────────┐

│               DataCleaningEnvironment (server/data_cleaning_env.py)  │

│                                                                      │

│  ┌─────────────┐  ┌──────────────┐  ┌───────────┐  ┌────────────┐ │

│  │ dataset_    │  │  Action      │  │  Grader   │  │  Reward    │ │

│  │ factory.py  │  │  Dispatcher  │  │  Engine   │  │  Computer  │ │

│  │             │  │  SET_VALUE   │  │  grade()  │  │            │ │

│  │ easy/medium │  │  DROP_ROW    │  │  score    │  │  progress  │ │

│  │ /hard CSVs  │  │  STANDARD.   │  │  delta    │  │  efficiency│ │

│  │             │  │  FILL_MISS.  │  │           │  │  penalties │ │

│  └─────────────┘  └──────────────┘  └───────────┘  └────────────┘ │

└─────────────────────────────────────────────────────────────────────┘

```

---

## 📁 Project Structure

```

data_cleaning_env/

│

├── 📄 client.py                  # DataCleaningEnv — OpenEnv client

├── 📄 models.py                  # CleanAction, CleanObservation, CleanState (Pydantic)

├── 📄 inference.py               # Official evaluation entry point

├── 📄 dataset_factory.py         # Generates easy/medium/hard dirty↔clean CSV pairs

├── 📄 graders.py                 # Scoring engine — grade(agent_df vs clean_df)

├── 📄 openenv.yaml               # OpenEnv manifest (HuggingFace Spaces config)

├── 📄 pyproject.toml             # Project metadata and dependencies

│

└── server/

    ├── 📄 app.py                 # FastAPI application (HTTP + WebSocket)

    ├── 📄 data_cleaning_env.py   # Core environment logic (reset/step/state)

    ├── 📄 __init__.py

    └── 📄 Dockerfile             # Container image definition

```

---

## 🎯 Tasks

The environment ships three progressively harder tasks, each with fixed-seed deterministic datasets:

### 🟢 Easy — Sales Orders

| Property | Value |
|---|---|
| Dataset | ~100-row sales orders CSV |
| Dirty Issues | Cell-level type errors, a few missing values |
| Step Budget | **40 steps** |
| Success Threshold | **Score ≥ 0.95** |
| Primary Skills | `SET_VALUE`, `FILL_MISSING` |

**What the agent needs to fix:** Individual cells with wrong types (e.g., `"N/A"` in a price column, `"abc"` in a numeric field). Straightforward injected errors with clear ground truth.

---

### 🟡 Medium — Financial Transactions

| Property | Value |
|---|---|
| Dataset | ~200-row transaction log |
| Dirty Issues | Outlier rows, mixed date formats, missing amounts |
| Step Budget | **80 steps** |
| Success Threshold | **Score ≥ 0.85** |
| Primary Skills | `DROP_ROW`, `STANDARDIZE_COL`, `FILL_MISSING` |

**What the agent needs to fix:** Statistical outliers disguised as data, inconsistent date formats, missing numeric values. Crucially, some extreme values are **valid** — dropping them costs a false-positive penalty.

---

### 🔴 Hard — Multi-Schema Dataset

| Property | Value |
|---|---|
| Dataset | ~400-row multi-domain CSV |
| Dirty Issues | Cross-column inconsistencies, future-year dates, bulk missing data |
| Step Budget | **150 steps** |
| Success Threshold | **Score ≥ 0.80** |
| Primary Skills | All commands |

**What the agent needs to fix:** Everything from easy + medium, plus cascading schema issues across columns. Requires strategic planning about fix order.

---

## 🕹️ Action Space

Every step the agent sends exactly one `CleanAction`:

```python

from models import CleanAction



# Fix a specific cell

CleanAction(command="SET_VALUE", row_index=3, column="price", value="29.99")



# Remove an entire row (use carefully — false positives are penalised)

CleanAction(command="DROP_ROW", row_index=17)



# Normalise a column's format (dates → YYYY-MM-DD, numbers → float, strings → stripped)

CleanAction(command="STANDARDIZE_COL", column="order_date")



# Fill all NaN values in a column using a strategy

CleanAction(command="FILL_MISSING", column="quantity", fill_strategy="median")



# Signal episode completion (only accepted when score ≥ task threshold)

CleanAction(command="DONE")

```

### Command Reference

| Command | `row_index` | `column` | `value` | `fill_strategy` |
|---|---|---|---|---|
| `SET_VALUE` | ✅ required | ✅ required | ✅ required | — |
| `DROP_ROW` | ✅ required | — | — | — |
| `STANDARDIZE_COL` | — | ✅ required | — | — |
| `FILL_MISSING` | — | ✅ required | — | ✅ required |
| `DONE` | — | — | — | — |

### `FILL_MISSING` Strategies



| Strategy | Behaviour |

|---|---|

| `"mean"` | Replace NaN with column mean (numeric columns only) |

| `"median"` | Replace NaN with column median (numeric columns only) |

| `"mode"` | Replace NaN with most frequent value (any column) |

| `"drop"` | Remove rows where this column is NaN |



> ⚠️ **Important:** `DROP_ROW` removes by **positional row index** (the `row_index` column in the CSV), not by a row ID field. Row indices shift after each drop.



---



## 👁️ Observation Space



After every `reset()` and `step()`, the agent receives a `CleanObservation`:



```python

@dataclass

class CleanObservation:

    # ── Task context (constant per episode) ──────────────────────

    task_id: str               # "easy" | "medium" | "hard"
    schema_hint: str           # Plain-English description of clean schema

    initial_dirty_cells: int   # Total dirty cells at episode start


    # ── Per-step state ───────────────────────────────────────────

    dirty_csv: str             # Full current CSV as string (all edits applied)

    current_score: float       # 0.0 → 1.0  (grader score vs ground truth)

    issues_remaining: int      # Approximate dirty cells still to fix

    step_number: int           # Steps taken so far

    max_steps: int             # Budget for this task


    # ── Last-action feedback ─────────────────────────────────────

    last_action_success: bool  # Whether previous action applied cleanly

    last_action_error: str     # Error message if success=False (else None)


    # ── Inherited ────────────────────────────────────────────────

    done: bool                 # True = episode ended

    reward: float | None       # Per-step reward (None after reset)

```


### Score Computation

The grader compares the agent's working DataFrame to the hidden ground-truth DataFrame:

```

score = (initial_dirty_cells - remaining_dirty_cells) / initial_dirty_cells

```

A score of `1.0` means perfect agreement with ground truth.

---

## 💰 Reward Function

The reward is dense and shaped to guide efficient, precise cleaning:

```

reward = progress_term

       + efficiency_bonus

       + false_positive_penalty

       + early_done_penalty

       + step_cost

```

| Component | Value | When |
|---|---|---|
| **Progress** | `current_score − previous_score` | Every step |
| **Efficiency bonus** | `+0.10 × (1 − steps_used/max_steps)` | Only when task is solved this step |
| **False-positive penalty** | `−0.15` | `DROP_ROW` removes a valid-extreme row (medium task) |
| **Early DONE penalty** | `−0.20` | `DONE` called with score < 0.60 |
| **Step cost** | `−0.005` | Every step (discourages padding) |
| **Premature DONE block** | `−1.00` | `DONE` below task threshold — episode *continues* |

**Reward range:** `[−0.5, +1.0]` (clipped)

### Termination Logic

The episode terminates when **any** of these is true:

1. ✅ `current_score >= task_threshold` (auto-terminated, efficiency bonus awarded)
2. ✅ Agent sends `DONE` and `current_score >= task_threshold` (accepted)
3. ⏱️ `step_count >= max_steps` (budget exhausted)

`DONE` is **refused** if the score is below threshold — the episode continues with a `−1.0` reward signal.

---

## 🚀 Quick Start

### Prerequisites

- Python 3.12+
- Docker Desktop (for containerised server)
- A free [HuggingFace token](https://huggingface.co/settings/tokens) (for the inference LLM)

### 1. Clone & Install

```bash

git clone https://github.com/Code-Knight-Debjit/Data-Cleaning-Environment.git

cd Data-Cleaning-Environment



# Create virtual environment

python -m venv .venv



# Activate (Windows PowerShell)

.venv\Scripts\Activate.ps1



# Activate (macOS/Linux)

source .venv/bin/activate



# Install dependencies

pip install -e .

```

### 2. Build the Docker Image

```bash

docker build -t openenv-data_cleaning:latest -f server/Dockerfile .

```

### 3. Set Your HuggingFace Token

```powershell

# Windows PowerShell

$env:HF_TOKEN = "hf_your_token_here"



# macOS / Linux

export HF_TOKEN="hf_your_token_here"

```

### 4. Run Inference

```bash

python inference.py

```

That's it! The script auto-starts the Docker container, runs the LLM agent through all three tasks (easy → medium → hard), and prints structured evaluation logs.

---

## 🤖 Running Inference

### Environment Variables

| Variable | Default | Description |
|---|---|---|
| `HF_TOKEN` | *(required)* | Your HuggingFace token for LLM API access |
| `API_BASE_URL` | `https://router.huggingface.co/v1` | LLM API endpoint |
| `MODEL_NAME` | `Qwen/Qwen2.5-72B-Instruct` | Model to use for inference |
| `LOCAL_IMAGE_NAME` | `openenv-data_cleaning:latest` | Docker image to launch |
| `ENV_BASE_URL` | `http://localhost:8000` | Direct server URL (if not using Docker) |

### Switching Models

```powershell

# Use Mistral (smaller, faster)

$env:MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"



# Use Llama

$env:MODEL_NAME = "meta-llama/Llama-3.1-8B-Instruct"

```

### Connecting to a Running Server (skip Docker)

```powershell

$env:LOCAL_IMAGE_NAME = ""   # must be empty string

$env:ENV_BASE_URL = "http://localhost:8000"

python inference.py

```

### Expected Output

```

API_BASE_URL     : https://router.huggingface.co/v1

MODEL_NAME       : Qwen/Qwen2.5-72B-Instruct

LOCAL_IMAGE_NAME : openenv-data_cleaning:latest

ENV_BASE_URL     : http://localhost:8000



[START] task=easy env=data_cleaning_env model=Qwen/Qwen2.5-72B-Instruct

[STEP]  step=1  action=FILL_MISSING  reward=0.12 done=false  error=null

[STEP]  step=2  action=SET_VALUE     reward=0.08 done=false  error=null

[STEP]  step=3  action=STANDARDIZE_COL reward=0.05 done=false error=null

...

[END]   success=true steps=18 score=0.97 rewards=0.12,0.08,...



[START] task=medium env=data_cleaning_env ...

...



════════════════════════════════════════════════════════

Task        Score    Reward  Steps  Pass

────────────────────────────────────────────────────────

easy       0.9712    1.3400     18   YES

medium     0.8823    2.1100     47   YES

hard       0.7640    1.8500     98    NO

════════════════════════════════════════════════════════

```

---

## 🔌 Environment API

### Using the Python Client Directly

```python

import asyncio

from client import DataCleaningEnv

from models import CleanAction



async def run():

    # Option A: Auto-start Docker container

    env = await DataCleaningEnv.from_docker_image("openenv-data_cleaning:latest")



    # Option B: Connect to an already-running server

    # env = DataCleaningEnv(base_url="http://localhost:8000")

    # await env.connect()



    try:

        # Reset for a specific task

        result = await env.reset(task_id="easy")

        obs = result.observation



        print(f"Score: {obs.current_score:.4f}")

        print(f"Issues: {obs.issues_remaining}")

        print(f"Schema: {obs.schema_hint}")



        # Take a step

        action = CleanAction(

            command="FILL_MISSING",

            column="price",

            fill_strategy="median"

        )

        result = await env.step(action)

        obs = result.observation



        print(f"Reward: {result.reward:.4f}")

        print(f"New score: {obs.current_score:.4f}")

        print(f"Action OK: {obs.last_action_success}")



        # Signal completion

        result = await env.step(CleanAction(command="DONE"))



    finally:

        await env.close()



asyncio.run(run())

```

### Using the Sync Wrapper

```python

from client import DataCleaningEnv

from models import CleanAction



env = DataCleaningEnv(base_url="http://localhost:8000").sync()



with env:

    result = env.reset(task_id="easy")

    result = env.step(CleanAction(command="STANDARDIZE_COL", column="order_date"))

    print(f"Score: {result.observation.current_score:.4f}")

```

### HTTP Endpoints

When the server is running, the following HTTP endpoints are available:

| Endpoint | Method | Description |
|---|---|---|
| `/health` | GET | Server health check |
| `/docs` | GET | Swagger / OpenAPI documentation |
| `/web` | GET | Interactive web UI |
| `/ws` | WebSocket | Persistent session endpoint |

---

## ⚙️ Configuration

### Step Budgets

```python

MAX_STEPS = {

    "easy":   40,

    "medium": 80,

    "hard":   150,

}

```

### Success Thresholds

```python

DONE_THRESHOLD = {

    "easy":   0.95,

    "medium": 0.85,

    "hard":   0.80,

}

```

### Reward Constants

| Constant | Value | Purpose |
|---|---|---|
| `STEP_COST` | `-0.005` | Per-step penalty to discourage padding |
| `EARLY_DONE_PENALTY` | `-0.20` | Penalty for `DONE` below score 0.60 |
| `EARLY_DONE_THRESHOLD` | `0.60` | Score floor for DONE without penalty |
| `FALSE_POSITIVE_PENALTY` | `-0.15` | Penalty for wrongly dropping a valid row |
| `EFFICIENCY_BONUS_WEIGHT` | `0.10` | Multiplier for early-completion bonus |

---

## ☁️ Deployment

### Deploy to HuggingFace Spaces

```bash

# Install the OpenEnv CLI

pip install openenv



# Authenticate with HuggingFace

huggingface-cli login



# Deploy (from the repo root where openenv.yaml lives)

openenv push



# Or deploy privately to a specific repo

openenv push --repo-id your-username/data-cleaning-env --private

```

After deployment, your environment will be live at:
```

https://huggingface.co/spaces/your-username/data-cleaning-env

```

With endpoints:
- **Web UI:** `/web`
- **API Docs:** `/docs`
- **Health:** `/health`
- **WebSocket:** `/ws`

### Connect to a HuggingFace Space

```python

env = await DataCleaningEnv.from_env("your-username/data-cleaning-env")

# or run locally with UV (no Docker needed)

env = await DataCleaningEnv.from_env("your-username/data-cleaning-env", use_docker=False)

```

### Run the Server Locally (Without Docker)

```bash

uvicorn server.app:app --reload --port 8000

```

---

## 🧪 Development & Testing

### Test the Environment Logic (No Server Needed)

```bash

# Runs a smoke test across all three tasks

python server/data_cleaning_env.py

```

Expected output:
```

────────────────────────────────────────────────────────────────

TASK: EASY

────────────────────────────────────────────────────────────────

reset()  → score=0.0000  issues=29  done=False

  CSV:  101 rows, 5 cols

  Hint: Sales orders dataset. price must be float...

step (bad col) → success=False  error='Column 'DOES_NOT_EXIST' not found...'

step (fix row=3 col='price') → success=True  score=0.0345  reward=0.0295

step (DONE, blocked)  → done=False  reward=-1.0  score=0.0345

...

All smoke tests passed.

```

### Test Pydantic Models

```bash

python models.py

```

### Test the Client Parser

```bash

python test_parse.py

```

### Run the Full Server Locally

```bash

uvicorn server.app:app --reload

# Open http://localhost:8000/docs for interactive API explorer

```

---

## 🔧 Troubleshooting

### `TypeError: Too few arguments for EnvClient`

**Cause:** Your `client.py` subclasses `EnvClient` with only 2 type parameters, but OpenEnv requires 3 (`ActT`, `ObsT`, `StateT`).

**Fix:**
```python

# ❌ Wrong

class DataCleaningEnv(EnvClient[CleanAction, CleanObservation]):



# ✅ Correct

class DataCleaningEnv(EnvClient[CleanAction, CleanObservation, dict]):

```

Also ensure `_parse_state` is implemented:
```python

def _parse_state(self, payload: dict) -> dict:

    return payload

```

---

### `ValidationError: Input should be 'SET_VALUE', 'DROP_ROW', ...`

**Cause:** Passing an invalid command string to `CleanAction`.

**Fix:** Only these 5 commands are valid:
```python

"SET_VALUE" | "DROP_ROW" | "STANDARDIZE_COL" | "FILL_MISSING" | "DONE"

```
There is no `"drop_column"` — columns cannot be dropped, only rows.

---

### `UnboundLocalError: cannot access local variable 'env'`

**Cause 1:** Docker image doesn't exist yet.
```bash

docker build -t openenv-data_cleaning:latest -f server/Dockerfile .

```

**Cause 2:** Stray test lines in `inference.py` referencing `env` before it's assigned.

**Fix:** Remove any manually added lines like `action = CleanAction(...)` or `result = await env.step(action)` from inside `main()`. The `main()` function should only call `run_episode()` — all action logic belongs inside that function.

---

### `DONE rejected: score X < required Y`

**This is expected behaviour, not a bug.** The environment refuses premature termination. The agent should continue cleaning until the score meets the task threshold.

---

### HuggingFace Router returns 401

Ensure your token is set:
```powershell

$env:HF_TOKEN = "hf_your_token_here"

```
Get a free token at [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens).

---

## 📐 Data Flow Diagram

```

                    ┌──────────────────────────────────┐

                    │   inference.py / custom agent    │

                    │                                  │

                    │  1. await env.reset(task_id=…)   │

                    │  2. obs = result.observation      │

                    │  3. build_prompt(obs) → LLM       │

                    │  4. parse_action(llm_output)      │

                    │  5. await env.step(action)        │

                    │  6. GOTO 2 until done             │

                    └──────────────┬───────────────────┘

                                   │

                    CleanAction (JSON over WebSocket)

                                   │

                                   ▼

                    ┌──────────────────────────────────┐

                    │        DataCleaningEnvironment    │

                    │                                  │

                    │  _apply_action()                 │

                    │    → mutates _dirty_df in-place  │

                    │                                  │

                    │  grade(agent_df vs clean_df)     │

                    │    → score ∈ [0.0, 1.0]          │

                    │                                  │

                    │  _compute_reward()               │

                    │    → progress + bonuses          │

                    │                                  │

                    │  _build_observation()            │

                    │    → CleanObservation            │

                    └──────────────────────────────────┘

```

---

## 🤝 Contributing

1. Fork the repository
2. Create a feature branch: `git checkout -b feature/my-improvement`
3. Run the smoke tests: `python server/data_cleaning_env.py`
4. Commit your changes: `git commit -m "feat: add my improvement"`
5. Push and open a Pull Request

---

## 📄 License

This project is licensed under the MIT License. See [LICENSE](LICENSE) for details.

---

<div align="center">

Built with ❤️ using [OpenEnv](https://github.com/meta-pytorch/OpenEnv) · [FastAPI](https://fastapi.tiangolo.com/) · [Pydantic](https://docs.pydantic.dev/) · [HuggingFace](https://huggingface.co/)

</div>