README.md

metadata

title: Dataclean Env
emoji: 🧹
colorFrom: blue
colorTo: green
sdk: docker
pinned: false
tags:
  - openenv
  - data-cleaning
  - etl
  - real-world
  - tabular
  - data-drift
  - streaming

DataClean OpenEnv

A real-world data cleaning and ETL environment for AI agent RL training, built on the OpenEnv framework by Meta-PyTorch and Hugging Face.

Agents learn to fix messy tabular data — filling nulls, normalising inconsistent values, removing duplicates, filtering outliers, performing multi-table ETL joins, and handling live data drift (novel: fresh dirty rows injected mid-episode every 5 steps).

---

Why This Environment

Data cleaning is one of the most common, time-consuming tasks in real data engineering. Every company with a data pipeline does this daily. Training an RL agent on this task has immediate real-world value — unlike game-based environments.

What makes this different from other OpenEnv submissions:

• Task 4 (Data Drift) is genuinely novel — no existing OpenEnv environment simulates live streaming row injection mid-episode
• Parallel baseline using concurrent.futures.ThreadPoolExecutor — all 4 tasks run simultaneously
• Deterministic graders using seeded dataset generation — perfectly reproducible scores
• Dense reward signal: reward = grader(new_state) − grader(prev_state) at every step

---

Tasks

Task	Difficulty	Description	Max Steps
task1	Easy	Fix nulls + dtypes in 50-row customer CSV	10
task2	Medium	Dedup + normalize strings + fix dates + fill nulls	20
task3	Hard	Multi-table merge + outlier removal + derived column	30
task4_data_drift	Expert	Live streaming table — 7 dirty rows injected every 5 steps	40

---

Action Space

Every action is a JSON object with an operation field:

{"operation": "fill_nulls",   "column": "age",          "strategy": "median"}
{"operation": "cast_column",  "column": "age",           "dtype": "int"}
{"operation": "remove_duplicates"}
{"operation": "normalize_values", "column": "country",   "method": "upper"}
{"operation": "cast_column",  "column": "order_date",    "dtype": "datetime"}
{"operation": "filter_outliers", "column": "amount",     "method": "iqr", "threshold": 1.5}
{"operation": "merge_tables", "left_table": "orders",    "right_table": "customers", "on": "customer_id", "output_table": "merged"}
{"operation": "add_derived_column", "column_name": "order_year", "source_column": "order_date", "transform": "year_from_date", "table_name": "merged"}
{"operation": "submit"}

Observation Space

After each reset() / step() the agent receives:

• task_id, task_description, step_count, max_steps, message
• tables — dict of {table_name → JSON string of df.head(10)}
• column_dtypes, null_counts, duplicate_count, row_count
• schema_errors — list of detected problems to guide the agent
• reward, done, partial_score — RL signals

---

Reward Function

step_reward = grader(current_state) − grader(previous_state)   # dense delta signal
invalid_op  = −0.02                                              # bad operation penalty
terminal    = final grader score on submit or max_steps

Partial credit per sub-dimension. Score range: [0.0, 1.0].

---

Baseline Scores (llama-3.3-70b-versatile, seed=42, parallel run)

Task	Score	Time
task1 (easy)	1.0000	11.8s
task2 (medium)	1.0000	34.1s
task3 (hard)	0.8000	22.4s
task4_data_drift (expert)	0.9297	26.6s
mean	0.9324	34.3s wall

All 4 tasks run in parallel — wall time = slowest task, not sum.

---

Setup & Local Run

pip install -r requirements.txt
uvicorn server.app:app --host 0.0.0.0 --port 7860 --reload

curl http://localhost:7860/health
curl http://localhost:7860/tasks

export OPENAI_API_KEY=your_key
export OPENAI_BASE_URL=https://api.groq.com/openai/v1
export BASELINE_MODEL=llama-3.3-70b-versatile
python baseline.py

Docker

docker build -t dataclean-env .
docker run -p 7860:7860 -e OPENAI_API_KEY=your_key dataclean-env

---

API Endpoints

Method	Endpoint	Description
POST	/reset	Start new episode
POST	/step	Execute one cleaning operation
GET	/state	Current episode metadata
GET	/tasks	All tasks + action schema
GET	/grader	Score current episode state
GET	/baseline	Run baseline agent on all tasks
GET	/health	Liveness probe
GET	/docs	Interactive Swagger UI

---

Project Structure

dataCleaningProject/
├── server/
│   ├── app.py              # FastAPI server — all endpoints
│   ├── environment.py      # Core env logic — reset/step/state + drift injection
│   ├── graders.py          # Deterministic scoring for all 4 tasks
│   └── dataset_factory.py  # Seeded dirty+expected dataset generation + drift batches
├── models.py               # Pydantic Action + Observation models
├── baseline.py             # Parallel baseline (ThreadPoolExecutor, 4 tasks at once)
├── client.py               # HTTP client
├── openenv.yaml            # OpenEnv manifest
├── Dockerfile              # Port 7860, Python 3.11-slim
└── requirements.txt