# 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 Cleaning Environment Implementation. Simulates real-world tabular data cleaning tasks where an AI agent must identify and fix data quality issues: missing values, duplicates, type errors, and outliers. Tasks (easy → medium → hard): 1. easy – Fix missing values in a small CSV (impute/drop) 2. medium – Remove duplicates AND fix type-cast errors 3. hard – Full pipeline: missing values + duplicates + outliers + normalise API: POST /reset – start a new episode (optionally pass {"task": "easy"|"medium"|"hard"}) POST /step – apply one cleaning operation GET /state – internal episode state GET /tasks – list tasks + action schema POST /grader – score a completed episode POST /baseline – run a deterministic baseline agent on all 3 tasks """ import copy import random import re from typing import Any, Optional from uuid import uuid4 from openenv.core.env_server.interfaces import Environment from openenv.core.env_server.types import State try: from ..models import DataCleaningAction, DataCleaningObservation except ImportError: from models import DataCleaningAction, DataCleaningObservation # --------------------------------------------------------------------------- # Synthetic datasets (deterministic, reproducible) # --------------------------------------------------------------------------- def _make_easy_dataset(seed: int = 0) -> list[dict]: """5-row CSV with 3 missing values to impute or drop.""" rng = random.Random(seed) names = ["Alice", None, "Charlie", "Diana", None] ages = [25, 30, None, 22, 28] rows = [] for n, a in zip(names, ages): rows.append({"name": n, "age": a, "score": round(rng.uniform(50, 100), 1)}) return rows def _make_medium_dataset(seed: int = 0) -> list[dict]: """8-row dataset with 2 duplicate rows and 2 type errors (age stored as str).""" base = [ {"id": 1, "name": "Alice", "age": 25, "city": "Delhi"}, {"id": 2, "name": "Bob", "age": "30x", "city": "Mumbai"}, # bad type {"id": 3, "name": "Charlie", "age": 22, "city": "Pune"}, {"id": 2, "name": "Bob", "age": "30x", "city": "Mumbai"}, # duplicate of row 2 {"id": 4, "name": "Diana", "age": "abc", "city": "Chennai"}, # bad type {"id": 5, "name": "Eve", "age": 27, "city": "Kolkata"}, {"id": 3, "name": "Charlie", "age": 22, "city": "Pune"}, # duplicate of row 3 {"id": 6, "name": "Frank", "age": 33, "city": "Hyderabad"}, ] return copy.deepcopy(base) def _make_hard_dataset(seed: int = 0) -> list[dict]: """12-row dataset: missing values + duplicates + outliers + unnormalised strings.""" rng = random.Random(seed) rows = [ {"id": 1, "product": " Widget A ", "price": 19.99, "quantity": None, "category": "Electronics"}, {"id": 2, "product": "Gadget B", "price": 9999.0, "quantity": 50, "category": "electronics"}, # outlier price + case {"id": 3, "product": "Widget A ", "price": 19.99, "quantity": None, "category": "Electronics"}, # dup of 1 {"id": 4, "product": "Doohickey C", "price": 5.49, "quantity": 200, "category": "HOME"}, {"id": 5, "product": None, "price": 12.00, "quantity": 75, "category": "Electronics"}, {"id": 6, "product": "Thingamajig", "price": -50.0, "quantity": 10, "category": "Home"}, # neg price outlier {"id": 7, "product": "Widget A", "price": 19.99, "quantity": 30, "category": "Electronics"}, # near-dup of 1 {"id": 8, "product": "Gizmo D", "price": 7.25, "quantity": None, "category": "Home"}, {"id": 9, "product": "Gadget B", "price": 9.99, "quantity": 50, "category": "Electronics"}, {"id": 10, "product": "Sprocket E", "price": 3.15, "quantity": 500, "category": " Electronics"}, {"id": 11, "product": "Cog F", "price": 1.99, "quantity": 1000, "category": "Home"}, {"id": 12, "product": "Bolt G", "price": 0.50, "quantity": None, "category": "home"}, ] return copy.deepcopy(rows) # --------------------------------------------------------------------------- # Grader helpers # --------------------------------------------------------------------------- def _grade_easy(rows: list[dict]) -> float: """Score 0-1: reward for each missing value resolved (imputed or dropped).""" if not rows: return 0.0 missing = sum(1 for r in rows for v in r.values() if v is None) total_possible = 3 # 2 missing names + 1 missing age in original resolved = max(0, total_possible - missing) return round(resolved / total_possible, 4) def _grade_medium(rows: list[dict]) -> float: """Score 0-1: half for no duplicates, half for no type errors in 'age'.""" if not rows: return 0.0 # Dedup check: unique (id, name) combos seen = set() dupes = 0 for r in rows: key = (r.get("id"), r.get("name")) if key in seen: dupes += 1 seen.add(key) dedup_score = 1.0 if dupes == 0 else max(0.0, 1.0 - dupes * 0.5) # Type-fix check: all 'age' values must be int or float type_errors = sum(1 for r in rows if not isinstance(r.get("age"), (int, float))) type_score = 1.0 if type_errors == 0 else max(0.0, 1.0 - type_errors * 0.5) return round((dedup_score + type_score) / 2, 4) def _grade_hard(rows: list[dict]) -> float: """Score 0-1 across 4 sub-criteria.""" if not rows: return 0.0 scores = [] # 1. No missing values missing = sum(1 for r in rows for v in r.values() if v is None) scores.append(1.0 if missing == 0 else max(0.0, 1.0 - missing * 0.1)) # 2. No duplicate products (after stripping) products = [str(r.get("product", "")).strip().lower() for r in rows if r.get("product")] dupes = len(products) - len(set(products)) scores.append(1.0 if dupes == 0 else max(0.0, 1.0 - dupes * 0.2)) # 3. No price outliers (price must be 0 < price < 500) outliers = sum(1 for r in rows if isinstance(r.get("price"), (int, float)) and not (0 < r["price"] < 500)) scores.append(1.0 if outliers == 0 else max(0.0, 1.0 - outliers * 0.3)) # 4. Normalised category strings (no leading/trailing spaces, title-case) bad_cats = sum( 1 for r in rows if isinstance(r.get("category"), str) and (r["category"] != r["category"].strip() or r["category"] != r["category"].title()) ) scores.append(1.0 if bad_cats == 0 else max(0.0, 1.0 - bad_cats * 0.15)) return round(sum(scores) / len(scores), 4) GRADERS = { "easy": _grade_easy, "medium": _grade_medium, "hard": _grade_hard, } DATASETS = { "easy": _make_easy_dataset, "medium": _make_medium_dataset, "hard": _make_hard_dataset, } # --------------------------------------------------------------------------- # Environment # --------------------------------------------------------------------------- MAX_STEPS = { "easy": 10, "medium": 15, "hard": 25, } VALID_OPERATIONS = { # Easy task ops "impute_mean", # fill numeric NaN with column mean "impute_mode", # fill categorical NaN with column mode "drop_missing_rows", # drop all rows that contain any None # Medium task ops "remove_duplicates", # drop exact duplicate rows "fix_type_errors", # coerce non-numeric 'age' / numeric cols to float (NaN if fails) # Hard task ops "remove_outliers", # drop rows where price < 0 or price > 500 "normalize_text", # strip + title-case all string columns "fill_quantity_mean", # fill missing quantity with column mean } class DataCleaningEnvironment(Environment): """ Data Cleaning RL Environment. The agent receives a dirty dataset as observation and must apply cleaning operations step-by-step to maximise data quality. Each task presents a different level of difficulty and requires different cleaning operations. """ SUPPORTS_CONCURRENT_SESSIONS: bool = True # ----------------------------------------------------------------------- # Lifecycle # ----------------------------------------------------------------------- def __init__(self): super().__init__() self._task: str = "easy" self._rows: list[dict] = [] self._original_rows: list[dict] = [] self._step_count: int = 0 self._episode_id: str = str(uuid4()) self._done: bool = False self._last_reward: float = 0.0 self._applied_ops: list[str] = [] def reset( self, seed: Optional[int] = None, episode_id: Optional[str] = None, task: str = "easy", **kwargs: Any, ) -> DataCleaningObservation: """ Reset the environment for a new episode. Args: seed: Random seed for dataset generation. episode_id: Optional custom episode identifier. task: One of "easy", "medium", "hard". """ if task not in DATASETS: task = "easy" self._task = task _seed = seed if seed is not None else random.randint(0, 9999) self._rows = DATASETS[task](_seed) self._original_rows = copy.deepcopy(self._rows) self._step_count = 0 self._episode_id = episode_id or str(uuid4()) self._done = False self._last_reward = 0.0 self._applied_ops = [] return self._make_observation(reward=0.0, done=False) def step( self, action: DataCleaningAction, timeout_s: Optional[float] = None, **kwargs: Any, ) -> DataCleaningObservation: """ Apply one cleaning operation to the dataset. Args: action: DataCleaningAction with fields: - operation (str): cleaning op name - column (str, optional): target column """ if self._done: return self._make_observation(reward=0.0, done=True) self._step_count += 1 op = action.operation.strip().lower() col = getattr(action, "column", None) if op not in VALID_OPERATIONS: # Invalid operation → small penalty, episode continues reward = -0.05 self._last_reward = reward obs = self._make_observation(reward=reward, done=False) obs.metadata["error"] = f"Unknown operation '{op}'. Valid: {sorted(VALID_OPERATIONS)}" return obs before_score = GRADERS[self._task](self._rows) self._apply_operation(op, col) after_score = GRADERS[self._task](self._rows) # Reward = improvement in quality score (partial progress signal) improvement = after_score - before_score if improvement > 0: reward = round(improvement + 0.02, 4) # bonus for any improvement elif improvement < 0: reward = round(improvement - 0.02, 4) # penalty for harming dataset else: reward = 0.0 # neutral if no change if after_score >= 1.0: reward += 1.0 # big bonus for completing the task self._last_reward = reward self._applied_ops.append(op) # Episode ends when perfect score or step limit reached max_steps = MAX_STEPS[self._task] done = (after_score >= 1.0) or (self._step_count >= max_steps) self._done = done return self._make_observation(reward=reward, done=done) # ----------------------------------------------------------------------- # State # ----------------------------------------------------------------------- @property def state(self) -> State: return State( episode_id=self._episode_id, step_count=self._step_count, task=self._task, done=self._done, current_score=GRADERS[self._task](self._rows), applied_ops=self._applied_ops, rows_remaining=len(self._rows), ) # ----------------------------------------------------------------------- # Grader (called via POST /grader) # ----------------------------------------------------------------------- def grade(self) -> dict: """Return final grader score for the current episode.""" raw = GRADERS[self._task](self._rows) score = max(0.001, min(0.999, raw)) # strictly (0, 1) as required return { "task": self._task, "score": score, "steps_taken": self._step_count, "ops_applied": self._applied_ops, } # ----------------------------------------------------------------------- # Tasks manifest (called via GET /tasks) # ----------------------------------------------------------------------- @staticmethod def tasks() -> list[dict]: return [ { "task_id": "easy", "description": "Fix missing values in a 5-row name/age/score table.", "difficulty": "easy", "max_steps": MAX_STEPS["easy"], "action_schema": { "operation": { "type": "string", "enum": ["impute_mean", "impute_mode", "drop_missing_rows"], }, "column": {"type": "string", "description": "Target column (optional)"}, }, }, { "task_id": "medium", "description": "Remove duplicate rows and fix type errors in the 'age' column.", "difficulty": "medium", "max_steps": MAX_STEPS["medium"], "action_schema": { "operation": { "type": "string", "enum": ["remove_duplicates", "fix_type_errors", "drop_missing_rows"], }, "column": {"type": "string", "description": "Target column (optional)"}, }, }, { "task_id": "hard", "description": ( "Full pipeline: fix missing values, remove duplicates, " "remove price outliers, and normalise text fields." ), "difficulty": "hard", "max_steps": MAX_STEPS["hard"], "action_schema": { "operation": { "type": "string", "enum": list(VALID_OPERATIONS), }, "column": {"type": "string", "description": "Target column (optional)"}, }, }, ] # ----------------------------------------------------------------------- # Baseline (called via POST /baseline) # ----------------------------------------------------------------------- def run_baseline(self) -> dict: """ Run a deterministic rule-based baseline agent on all 3 tasks. Returns scores dict compatible with the hackathon /baseline endpoint. """ results = {} for task_id in ["easy", "medium", "hard"]: self.reset(seed=42, task=task_id) ops = _BASELINE_POLICIES[task_id] for op in ops: if not self._done: self.step(DataCleaningAction(operation=op)) raw = GRADERS[task_id](self._rows) results[task_id] = max(0.001, min(0.999, raw)) # strictly (0, 1) return {"baseline_scores": results} # ----------------------------------------------------------------------- # Internal helpers # ----------------------------------------------------------------------- def _make_observation(self, reward: float, done: bool) -> DataCleaningObservation: missing_count = sum(1 for r in self._rows for v in r.values() if v is None) has_dupes = self._has_duplicates() has_outliers = self._has_outliers() score = GRADERS[self._task](self._rows) return DataCleaningObservation( current_text=self._rows_to_text(), is_normalized=not has_outliers and missing_count == 0 and not has_dupes, html_found=False, remaining_typos=missing_count + (2 if has_dupes else 0) + (1 if has_outliers else 0), done=done, reward=reward, metadata={ "task": self._task, "step": self._step_count, "rows": copy.deepcopy(self._rows), "missing_count": missing_count, "has_duplicates": has_dupes, "has_outliers": has_outliers, "quality_score": score, "valid_operations": sorted(VALID_OPERATIONS), "ops_already_applied": list(self._applied_ops), "recommended_next": _recommend_next( self._task, missing_count, has_dupes, has_outliers, self._applied_ops ), }, ) def _rows_to_text(self) -> str: if not self._rows: return "[]" cols = list(self._rows[0].keys()) header = " | ".join(cols) lines = [header, "-" * len(header)] for r in self._rows: lines.append(" | ".join(str(r.get(c, "")) for c in cols)) return "\n".join(lines) def _has_duplicates(self) -> bool: seen = set() for r in self._rows: key = tuple( (k, str(v).strip().lower() if isinstance(v, str) else v) for k, v in sorted(r.items()) ) if key in seen: return True seen.add(key) return False def _has_outliers(self) -> bool: return any( isinstance(r.get("price"), (int, float)) and not (0 < r["price"] < 500) for r in self._rows ) def _apply_operation(self, op: str, column: Optional[str]): """Mutate self._rows according to the chosen operation.""" if op == "drop_missing_rows": self._rows = [r for r in self._rows if all(v is not None for v in r.values())] elif op == "impute_mean": cols = [column] if column else _numeric_cols(self._rows) for c in cols: vals = [r[c] for r in self._rows if isinstance(r.get(c), (int, float))] if vals: mean = sum(vals) / len(vals) for r in self._rows: if r.get(c) is None: r[c] = round(mean, 2) elif op == "impute_mode": cols = [column] if column else _string_cols(self._rows) for c in cols: vals = [r[c] for r in self._rows if r.get(c) is not None] if vals: mode = max(set(vals), key=vals.count) for r in self._rows: if r.get(c) is None: r[c] = mode elif op == "remove_duplicates": seen: set = set() unique = [] for r in self._rows: key = tuple(sorted(r.items())) if key not in seen: seen.add(key) unique.append(r) self._rows = unique elif op == "fix_type_errors": cols = [column] if column else _numeric_cols(self._rows) + ["age"] for c in set(cols): for r in self._rows: val = r.get(c) if val is not None and not isinstance(val, (int, float)): # Try to coerce to float; set None if it fails try: r[c] = float(re.sub(r"[^\d.\-]", "", str(val))) except ValueError: r[c] = None elif op == "remove_outliers": self._rows = [ r for r in self._rows if not isinstance(r.get("price"), (int, float)) or (0 < r["price"] < 500) ] elif op == "normalize_text": for r in self._rows: for k, v in list(r.items()): if isinstance(v, str): r[k] = v.strip().title() elif op == "fill_quantity_mean": vals = [r["quantity"] for r in self._rows if isinstance(r.get("quantity"), (int, float))] if vals: mean = round(sum(vals) / len(vals), 1) for r in self._rows: if r.get("quantity") is None: r["quantity"] = mean # --------------------------------------------------------------------------- # Baseline policies (deterministic rule-based agents) # --------------------------------------------------------------------------- _BASELINE_POLICIES: dict[str, list[str]] = { "easy": ["impute_mean", "impute_mode", "drop_missing_rows"], "medium": ["remove_duplicates", "fix_type_errors", "drop_missing_rows"], "hard": [ "drop_missing_rows", "fill_quantity_mean", "remove_duplicates", "fix_type_errors", "remove_outliers", "normalize_text", ], } # --------------------------------------------------------------------------- # Recommendation helper (guides the LLM toward correct next op) # --------------------------------------------------------------------------- def _recommend_next( task: str, missing_count: int, has_dupes: bool, has_outliers: bool, applied_ops: list, ) -> str: """Return a plain-English hint for the LLM about the best next operation.""" applied = set(applied_ops) if task == "easy": if missing_count > 0: return "There are missing values. Use impute_mean (numeric) or impute_mode (text)." return "No issues remain. Episode should be complete." if task == "medium": if has_dupes and "remove_duplicates" not in applied: return "Duplicate rows exist. Use remove_duplicates." if "fix_type_errors" not in applied: return "Non-numeric values in numeric columns. Use fix_type_errors." if missing_count > 0 and "drop_missing_rows" not in applied: return "Some values still missing after type fix. Use drop_missing_rows." return "No issues remain. Episode should be complete." # hard if missing_count > 0 and "fill_quantity_mean" not in applied: return "Missing quantity values. Use fill_quantity_mean first." if missing_count > 0 and "drop_missing_rows" not in applied: return "Missing product/category values. Use drop_missing_rows." if has_outliers and "remove_outliers" not in applied: return "Price outliers detected (price<=0 or price>=500). Use remove_outliers." if "normalize_text" not in applied: return "String columns have inconsistent casing/whitespace. Use normalize_text." if has_dupes and "remove_duplicates" not in applied: return "Duplicate rows remain. Use remove_duplicates." return "All issues fixed. Episode should be complete." \ # --------------------------------------------------------------------------- # Column utility helpers # --------------------------------------------------------------------------- def _numeric_cols(rows: list[dict]) -> list[str]: if not rows: return [] return [ c for c in rows[0] if any(isinstance(r.get(c), (int, float)) for r in rows) ] def _string_cols(rows: list[dict]) -> list[str]: if not rows: return [] return [ c for c in rows[0] if any(isinstance(r.get(c), str) for r in rows) ]