| from __future__ import annotations
|
|
|
| from typing import Optional, Tuple
|
|
|
| import numpy as np
|
| import pandas as pd
|
|
|
| from .graders import grade_task
|
| from .models import Action, Observation, StepResult, TablePreview
|
| from .rewards import compute_reward
|
| from .tasks import TASK_IDS, get_task
|
|
|
|
|
| class DataCleaningEnv:
|
| MAX_STEPS: int = 20
|
| MIN_EPISODE_SCORE: float = 0.01
|
| MAX_EPISODE_SCORE: float = 0.99
|
|
|
| def __init__(self) -> None:
|
| self.task_id: Optional[str] = None
|
| self._task_config: Optional[dict] = None
|
| self.original_df: Optional[pd.DataFrame] = None
|
| self.current_df: Optional[pd.DataFrame] = None
|
| self.step_count: int = 0
|
| self.cleaning_log: list = []
|
| self.action_history: list = []
|
| self.raw_cumulative_reward: float = 0.0
|
| self.cumulative_reward: float = 0.0
|
| self.done: bool = False
|
| self.final_score: float = 0.01
|
|
|
| def reset(self, task_id: Optional[str] = None) -> Observation:
|
| if task_id is None:
|
| task_id = TASK_IDS[0]
|
| self.task_id = task_id
|
| self._task_config = get_task(task_id)
|
| self.original_df = self._task_config["dirty_df"].copy()
|
| self.current_df = self._task_config["dirty_df"].copy()
|
| self.step_count = 0
|
| self.cleaning_log = []
|
| self.action_history = []
|
| self.raw_cumulative_reward = 0.0
|
| self.cumulative_reward = 0.0
|
| self.done = False
|
| self.final_score = 0.01
|
| return self._build_observation()
|
|
|
| def step(self, action: Action) -> StepResult:
|
| if self.done:
|
| return StepResult(
|
| observation=self._build_observation(),
|
| reward=self.final_score,
|
| done=True,
|
| info={
|
| "error": "Episode already finished",
|
| "cumulative_reward": self.cumulative_reward,
|
| "raw_cumulative_reward": self.raw_cumulative_reward,
|
| "final_score": self.final_score,
|
| "step": self.step_count,
|
| },
|
| )
|
|
|
| error: Optional[str] = None
|
| reward: float = 0.0
|
|
|
| if action.type == "submit":
|
| self.final_score = grade_task(self.task_id, self.current_df)
|
| reward = self.final_score
|
| self.cleaning_log.append(f"[SUBMIT] Final grade: {self.final_score:.4f}")
|
| self.done = True
|
| else:
|
| try:
|
| reward, log_msg = self._apply_action(action)
|
| self.cleaning_log.append(log_msg)
|
| except Exception as exc:
|
| error = str(exc)
|
| reward = -0.10
|
| self.cleaning_log.append(f"[ERROR] {error}")
|
|
|
| self.step_count += 1
|
| self.raw_cumulative_reward = round(self.raw_cumulative_reward + reward, 4)
|
| self.cumulative_reward = self._clamp_episode_score(self.raw_cumulative_reward)
|
| self.action_history.append(action.model_dump())
|
|
|
| if not self.done and self.step_count >= self.MAX_STEPS:
|
| self.final_score = grade_task(self.task_id, self.current_df)
|
| self.done = True
|
|
|
| return StepResult(
|
| observation=self._build_observation(),
|
| reward=round(reward, 4),
|
| done=self.done,
|
| info={
|
| "error": error,
|
| "cumulative_reward": self.cumulative_reward,
|
| "raw_cumulative_reward": self.raw_cumulative_reward,
|
| "final_score": self.final_score,
|
| "step": self.step_count,
|
| },
|
| )
|
|
|
| def state(self) -> dict:
|
| return {
|
| "task_id": self.task_id,
|
| "step_count": self.step_count,
|
| "cumulative_reward": self.cumulative_reward,
|
| "raw_cumulative_reward": self.raw_cumulative_reward,
|
| "final_score": self.final_score,
|
| "done": self.done,
|
| "cleaning_log": self.cleaning_log,
|
| "action_history": self.action_history,
|
| "current_data": self._df_records_with_none(self.current_df) if self.current_df is not None else [],
|
| }
|
|
|
| @classmethod
|
| def _clamp_episode_score(cls, value: float) -> float:
|
| return round(min(max(value, cls.MIN_EPISODE_SCORE), cls.MAX_EPISODE_SCORE), 4)
|
|
|
| def _apply_action(self, action: Action) -> Tuple[float, str]:
|
| df = self.current_df
|
|
|
| if action.type == "fill_missing":
|
| col = self._require_column(action.column, df)
|
| missing_before = int(df[col].isna().sum())
|
| if missing_before == 0:
|
| return -0.05, f"[WARN] No missing values in '{col}' — wasted step"
|
|
|
| if action.strategy == "mean":
|
| df[col] = df[col].fillna(df[col].mean())
|
| elif action.strategy == "median":
|
| df[col] = df[col].fillna(df[col].median())
|
| elif action.strategy == "mode":
|
| df[col] = df[col].fillna(df[col].mode().iloc[0])
|
| elif action.strategy == "constant":
|
| df[col] = df[col].fillna(action.value)
|
| else:
|
| raise ValueError(f"Unknown fill strategy '{action.strategy}'")
|
|
|
| reward = compute_reward("fill_missing", {"filled": missing_before})
|
| return reward, f"Filled {missing_before} missing values in '{col}' via {action.strategy}"
|
|
|
| if action.type == "standardize_values":
|
| col = self._require_column(action.column, df)
|
| if not action.mapping:
|
| raise ValueError("'mapping' dict is required for standardize_values")
|
| replaced = int(df[col].isin(action.mapping.keys()).sum())
|
| df[col] = df[col].apply(lambda x: action.mapping.get(str(x), x) if pd.notna(x) else x)
|
| reward = compute_reward("standardize_values", {"replaced": replaced})
|
| return reward, f"Standardised {replaced} values in '{col}'"
|
|
|
| if action.type == "remove_duplicates":
|
| before = len(df)
|
| self.current_df = df.drop_duplicates().reset_index(drop=True)
|
| removed = before - len(self.current_df)
|
| if removed == 0:
|
| return -0.05, "[WARN] No exact duplicates found — wasted step"
|
| reward = compute_reward("remove_duplicates", {"removed": removed})
|
| return reward, f"Removed {removed} duplicate row(s)"
|
|
|
| if action.type == "remove_row":
|
| if action.row_id is None:
|
| raise ValueError("'row_id' is required for remove_row")
|
| if action.row_id not in df.index:
|
| raise ValueError(f"Row index {action.row_id} not found (valid range 0–{len(df)-1})")
|
| self.current_df = df.drop(index=action.row_id).reset_index(drop=True)
|
| reward = compute_reward("remove_row", {})
|
| return reward, f"Removed row at index {action.row_id}"
|
|
|
| if action.type == "convert_type":
|
| col = self._require_column(action.column, df)
|
| tgt = action.target_type
|
|
|
| if tgt == "float":
|
| df[col] = (
|
| df[col]
|
| .astype(str)
|
| .str.replace(r"[$,\s]", "", regex=True)
|
| .replace("nan", np.nan)
|
| .replace("None", np.nan)
|
| )
|
| df[col] = pd.to_numeric(df[col], errors="coerce")
|
| elif tgt == "int":
|
| df[col] = pd.to_numeric(df[col], errors="coerce").astype("Int64")
|
| elif tgt == "str":
|
| df[col] = df[col].astype(str)
|
| elif tgt == "datetime":
|
| parsed = pd.to_datetime(df[col], errors="coerce")
|
| df[col] = parsed.dt.strftime("%Y-%m-%d")
|
| else:
|
| raise ValueError(f"Unknown target_type '{tgt}'")
|
|
|
| reward = compute_reward("convert_type", {})
|
| return reward, f"Converted column '{col}' → {tgt}"
|
|
|
| if action.type == "clip_outliers":
|
| col = self._require_column(action.column, df)
|
| if action.lower is None and action.upper is None:
|
| raise ValueError("At least one of 'lower' or 'upper' must be set")
|
|
|
| series = pd.to_numeric(df[col], errors="coerce")
|
| clipped = 0
|
| if action.lower is not None:
|
| clipped += int((series < action.lower).sum())
|
| if action.upper is not None:
|
| clipped += int((series > action.upper).sum())
|
|
|
| df[col] = series.clip(lower=action.lower, upper=action.upper)
|
| reward = compute_reward("clip_outliers", {"clipped": clipped})
|
| return reward, f"Clipped '{col}' to [{action.lower}, {action.upper}] ({clipped} value(s) affected)"
|
|
|
| raise ValueError(f"Unknown action type '{action.type}'")
|
|
|
| @staticmethod
|
| def _require_column(col: Optional[str], df: pd.DataFrame) -> str:
|
| if not col:
|
| raise ValueError("'column' field is required for this action")
|
| if col not in df.columns:
|
| raise ValueError(f"Column '{col}' not found. Available: {list(df.columns)}")
|
| return col
|
|
|
| @staticmethod
|
| def _df_records_with_none(df: pd.DataFrame) -> list[dict]:
|
| safe_df = df.astype(object).where(pd.notna(df), None)
|
| return safe_df.to_dict(orient="records")
|
|
|
| def _build_observation(self) -> Observation:
|
| df = self.current_df
|
| issues: list = []
|
|
|
| if df is not None:
|
| for col in df.columns:
|
| miss = int(df[col].isna().sum())
|
| if miss > 0:
|
| issues.append(f"Column '{col}' has {miss} missing value(s)")
|
| dup = int(df.duplicated().sum())
|
| if dup > 0:
|
| issues.append(f"{dup} exact duplicate row(s) detected")
|
|
|
| head = df.head(10).copy()
|
| head.insert(0, "_row_id", head.index.tolist())
|
| preview_rows = self._df_records_with_none(head)
|
| schema_info = {c: str(df[c].dtype) for c in df.columns}
|
| shape = list(df.shape)
|
| else:
|
| preview_rows, schema_info, shape = [], {}, [0, 0]
|
|
|
| preview = TablePreview(
|
| columns=["_row_id"] + (list(df.columns) if df is not None else []),
|
| rows=preview_rows,
|
| shape=shape,
|
| )
|
|
|
| return Observation(
|
| task_id=self.task_id or "",
|
| task_description=(self._task_config["description"] if self._task_config else ""),
|
| table_preview=preview,
|
| schema_info=schema_info,
|
| valid_actions=[
|
| "fill_missing",
|
| "standardize_values",
|
| "remove_duplicates",
|
| "remove_row",
|
| "convert_type",
|
| "clip_outliers",
|
| "submit",
|
| ],
|
| step=self.step_count,
|
| max_steps=self.MAX_STEPS,
|
| cleaning_log=self.cleaning_log[-6:],
|
| issues_detected=issues,
|
| )
|
|
|
