SimonPaul06/data-cleaning-env

1"""2DataCleaningEnv — OpenEnv-compliant environment3step() / reset() / state() API4"""5from __future__ import annotations6 7import json8from typing import Any, Dict, List, Optional, Tuple9 10import pandas as pd11from pydantic import BaseModel, Field12 13from tasks import TASKS14 15 16# ─────────────────────────────────────────────17# Typed Pydantic models (OpenEnv spec)18# ─────────────────────────────────────────────19 20class Observation(BaseModel):21    task_id: str22    task_description: str23    difficulty: str24    dataframe_json: str = Field(description="Current DataFrame as JSON string")25    columns: List[str]26    shape: List[int]27    null_counts: Dict[str, int]28    step: int29    max_steps: int30    last_action_result: str31    done: bool32 33 34class Action(BaseModel):35    action_type: str = Field(36        description=(37            "One of: fix_null | fix_dtype | deduplicate | standardize | "38            "drop_column | replace_value | noop"39        )40    )41    column: Optional[str] = Field(default=None, description="Target column name. For deduplicate, use as subset key (e.g. 'customer_id')")42    strategy: Optional[str] = Field(43        default=None,44        description="For fix_null: mean | median | mode | drop | value",45    )46    value: Optional[str] = Field(47        default=None,48        description="For fix_null with strategy='value', or replace_value old value",49    )50    new_value: Optional[str] = Field(51        default=None, description="For replace_value: new value to set"52    )53    dtype: Optional[str] = Field(54        default=None, description="For fix_dtype: int | float | str"55    )56    fmt: Optional[str] = Field(57        default=None,58        description="For standardize: date | phone | lowercase | uppercase",59    )60 61 62class StepResult(BaseModel):63    observation: Observation64    reward: float65    done: bool66    info: Dict[str, Any]67 68 69# ─────────────────────────────────────────────70# Environment71# ─────────────────────────────────────────────72 73class DataCleaningEnv:74    """75    OpenEnv-compliant data cleaning environment.76 77    Usage:78        env = DataCleaningEnv(task_id="easy")79        obs = env.reset()80        result = env.step(Action(action_type="deduplicate"))81        state = env.state()82    """83 84    def __init__(self, task_id: str = "easy") -> None:85        if task_id not in TASKS:86            raise ValueError(f"Unknown task_id '{task_id}'. Choose from: {list(TASKS)}")87        self.task_id = task_id88        self.task = TASKS[task_id]89        self._df: Optional[pd.DataFrame] = None90        self._step_count: int = 091        self._done: bool = False92        self._last_action_result: str = "Not started."93 94    # ── Public API ──────────────────────────────95 96    def reset(self) -> Observation:97        """Reset environment to initial state and return first observation."""98        self._df = self.task["get_data"]().copy()99        self._step_count = 0100        self._done = False101        self._last_action_result = "Environment reset. DataFrame loaded."102        return self._observe()103 104    def step(self, action: Action) -> StepResult:105        """Apply an action, compute reward, return StepResult."""106        if self._df is None:107            raise RuntimeError("Call reset() before step().")108 109        if self._done:110            return StepResult(111                observation=self._observe(),112                reward=self._grade(),113                done=True,114                info={"warning": "Episode already finished."},115            )116 117        self._step_count += 1118        result_msg = self._apply_action(action)119        self._last_action_result = result_msg120 121        reward = self._grade()122        max_steps = self.task["max_steps"]123        if reward >= 1.0 or self._step_count >= max_steps:124            self._done = True125 126        return StepResult(127            observation=self._observe(),128            reward=reward,129            done=self._done,130            info={"action_result": result_msg, "reward": reward},131        )132 133    def state(self) -> Dict[str, Any]:134        """Return serializable snapshot of current environment state."""135        import math136        def safe_dict(d):137            """Recursively replace NaN/Inf with None for JSON safety."""138            if isinstance(d, dict):139                return {k: safe_dict(v) for k, v in d.items()}140            elif isinstance(d, float) and (math.isnan(d) or math.isinf(d)):141                return None142            elif isinstance(d, list):143                return [safe_dict(i) for i in d]144            return d145 146        return {147            "task_id": self.task_id,148            "difficulty": self.task["difficulty"],149            "step": self._step_count,150            "max_steps": self.task["max_steps"],151            "done": self._done,152            "current_reward": self._grade() if self._df is not None else 0.0,153            "dataframe": safe_dict(self._df.where(self._df.notna(), None).to_dict()) if self._df is not None else {},154        }155 156    # ── Internal helpers ─────────────────────────157 158    def _grade(self) -> float:159        if self._df is None:160            return 0.0161        return self.task["grader"](self._df)162 163    def _observe(self) -> Observation:164        df = self._df165        null_counts: Dict[str, int] = {}166        df_json = "{}"167        columns: List[str] = []168        shape: List[int] = [0, 0]169 170        if df is not None:171            null_counts = {k: int(v) for k, v in df.isnull().sum().items()}172            df_json = df.head(20).to_json()  # limit payload size173            columns = list(df.columns)174            shape = list(df.shape)175 176        return Observation(177            task_id=self.task_id,178            task_description=self.task["description"],179            difficulty=self.task["difficulty"],180            dataframe_json=df_json,181            columns=columns,182            shape=shape,183            null_counts=null_counts,184            step=self._step_count,185            max_steps=self.task["max_steps"],186            last_action_result=self._last_action_result,187            done=self._done,188        )189 190    def _apply_action(self, action: Action) -> str:  # noqa: C901191        df = self._df192        try:193            atype = action.action_type.lower().strip()194 195            if atype == "noop":196                return "No operation."197 198            elif atype == "fix_null":199                col = action.column200                if col not in df.columns:201                    return f"ERROR: column '{col}' not found."202                before = int(df[col].isnull().sum())203                if before == 0:204                    return f"No nulls in '{col}'."205                strat = (action.strategy or "").lower()206                if strat == "mean":207                    df[col] = df[col].fillna(df[col].mean())208                elif strat == "median":209                    df[col] = df[col].fillna(df[col].median())210                elif strat == "mode":211                    df[col] = df[col].fillna(df[col].mode()[0])212                elif strat == "drop":213                    df.dropna(subset=[col], inplace=True)214                    df.reset_index(drop=True, inplace=True)215                elif strat == "value" and action.value is not None:216                    df[col] = df[col].fillna(action.value)217                else:218                    return f"ERROR: unknown strategy '{strat}'."219                after = int(df[col].isnull().sum())220                return f"fix_null '{col}' ({strat}): {before} → {after} nulls."221 222            elif atype == "fix_dtype":223                col = action.column224                if col not in df.columns:225                    return f"ERROR: column '{col}' not found."226                dtype_map = {"int": "Int64", "float": "float64", "str": "string"}227                target = dtype_map.get((action.dtype or "").lower(), action.dtype)228                try:229                    df[col] = pd.to_numeric(df[col], errors="coerce") if target in ("Int64", "float64") else df[col].astype(target)230                    if target == "Int64":231                        df[col] = df[col].round().astype("Int64")232                    return f"fix_dtype '{col}' → {target}."233                except Exception as e:234                    return f"ERROR converting '{col}': {e}"235 236            elif atype == "deduplicate":237                before = len(df)238                subset = [action.column] if action.column and action.column in df.columns else None239                df.drop_duplicates(subset=subset, inplace=True)240                df.reset_index(drop=True, inplace=True)241                self._df = df242                subset_label = f" by '{action.column}'" if subset else ""243                return f"deduplicate{subset_label}: {before} → {len(df)} rows (removed {before - len(df)})."244 245            elif atype == "standardize":246                col = action.column247                if col not in df.columns:248                    return f"ERROR: column '{col}' not found."249                fmt = (action.fmt or "").lower()250                if fmt == "date":251                    df[col] = pd.to_datetime(df[col], errors="coerce").dt.strftime("%Y-%m-%d")252                    return f"standardize '{col}' → date YYYY-MM-DD."253                elif fmt == "phone":254                    df[col] = df[col].astype(str).str.replace(r"\D", "", regex=True)255                    return f"standardize '{col}' → digits only."256                elif fmt == "lowercase":257                    df[col] = df[col].astype(str).str.lower().str.strip()258                    return f"standardize '{col}' → lowercase."259                elif fmt == "uppercase":260                    df[col] = df[col].astype(str).str.upper().str.strip()261                    return f"standardize '{col}' → uppercase."262                else:263                    return f"ERROR: unknown fmt '{fmt}'."264 265            elif atype == "drop_column":266                col = action.column267                if col not in df.columns:268                    return f"ERROR: column '{col}' not found."269                df.drop(columns=[col], inplace=True)270                return f"drop_column '{col}'."271 272            elif atype == "replace_value":273                col = action.column274                if col not in df.columns:275                    return f"ERROR: column '{col}' not found."276                old_val = action.value277                new_val = action.new_value278                count = int((df[col] == old_val).sum())279                df[col] = df[col].replace(old_val, new_val)280                return f"replace_value '{col}': '{old_val}' → '{new_val}' ({count} replaced)."281 282            else:283                return f"ERROR: unknown action_type '{atype}'."284 285        except Exception as exc:286            return f"ERROR: {exc}"287