Sm60786/data-cleaning-env

Data Cleaning Environment for OpenEnv

An OpenEnv-compliant RL environment where AI agents learn to clean messy, real-world datasets through structured text commands. The agent observes data quality issues and iteratively applies cleaning operations — scored against a deterministic ground truth.

Why data cleaning? Data professionals spend 60-80% of their time wrangling dirty data. This environment lets you train and evaluate agents on a task that has enormous real-world value, with clear success criteria and rich partial-credit rewards.

Quick Start

1. Install

pip install -e .

2. Run the server

# Option A: directly
python -m server.app

# Option B: via Docker
docker build -t data-cleaning-env .
docker run -p 7860:7860 data-cleaning-env

3. Run the baseline agent

export API_BASE_URL="https://api.openai.com/v1"
export MODEL_NAME="gpt-4o"
export OPENAI_API_KEY="sk-..."
export ENV_URL="http://localhost:7860"

python inference.py

Environment Description

The agent receives a dirty dataset (a pandas DataFrame) and must clean it by issuing one command per step. After each step the environment returns:

• —A data preview (first/last rows)
• —A data summary (column names, types, null counts)
• —A quality report (scored 0.0–1.0 across four dimensions)
• —The result or error of the last command

The episode ends when the agent issues submit or exceeds 30 steps.

Reward Function

Each step reward = delta(quality_score): positive when the data improves, negative when it degrades. A small penalty (-0.01) is applied for no-ops and errors.

The quality score combines four equally-weighted dimensions:

Action Space

The agent sends a DataCleaningAction(command="..."). The following commands are supported:

Observation Space

DataCleaningObservation fields:

Tasks

Task 1 — Easy: Customer Contact Cleanup

• —Size: ~50 rows, 6 columns (id, name, email, phone, city, signup_date)
• —Issues: inconsistent name casing, extra whitespace, mixed date formats, inconsistent city casing, duplicate rows
• —Expected difficulty: Straightforward — an agent that trims whitespace, standardizes case, and drops duplicates should score > 0.8

Task 2 — Medium: Sales Transaction Cleaning

• —Size: ~100 rows, 8 columns (transactionid, date, productname, category, quantity, unitprice, customerid, region)
• —Issues: missing values in multiple columns, wrong dtypes (strings instead of numbers), negative quantities, inconsistent product names, inconsistent region casing, duplicate transactions
• —Expected difficulty: Requires type casting, filling/dropping missing values, and value standardization

Task 3 — Hard: Employee HR Data Reconciliation

• —Size: ~120 rows, 10 columns (employeeid, firstname, lastname, email, department, hiredate, salary, managerid, performancerating, status)
• —Issues: department codes instead of names, inconsistent status values, name casing/whitespace, mixed date formats, invalid emails, salary outliers, out-of-range performance ratings, invalid manager references, missing values, duplicate records from merged systems
• —Expected difficulty: Requires multi-step reasoning, cross-field validation, and value mapping

API Endpoints

The server exposes the standard OpenEnv HTTP API:

Baseline Scores

Approximate scores using GPT-4o with the provided inference.py:

(Scores may vary slightly depending on LLM temperature and API version.)

Project Structure

├── README.md                 # This file
├── openenv.yaml              # OpenEnv manifest
├── pyproject.toml            # Python package config
├── Dockerfile                # Container for HF Spaces
├── inference.py              # Baseline inference script
├── __init__.py               # Package exports
├── models.py                 # Pydantic action/observation/state models
├── client.py                 # EnvClient subclass
├── server/
│   ├── __init__.py
│   ├── app.py                # FastAPI app (create_app + main)
│   └── data_cleaning_environment.py  # Core environment logic
├── tasks/
│   ├── __init__.py
│   └── definitions.py        # 3 task definitions with data generators
└── graders/
    ├── __init__.py
    └── grader.py              # Deterministic scoring (4 dimensions)

Environment Variables for Inference

License

MIT