"""Core KaggleSimEnv environment v3. Causal logic, hierarchical categories, failure-mode traps, contextual rewards. """ from __future__ import annotations from typing import Any from kaggle_sim_env.hints import HintProvider from kaggle_sim_env.leaderboard import Leaderboard from kaggle_sim_env.models import ( Action, ActionType, EnvState, FailureMode, Observation, Reward, RewardBreakdown, StepResponse, get_allowed_values, get_param_key, validate_category, ) from kaggle_sim_env.rewards import compute_reward from kaggle_sim_env.tasks import TaskDefinition, get_task class KaggleSimEnv: """RL environment simulating a Kaggle competition with causal dataset logic.""" def __init__(self) -> None: self._task: TaskDefinition | None = None self._cv_score: float = 0.0 self._test_score: float = 0.0 self._applied: list[str] = [] self._history: list[str] = [] self._step_count: int = 0 self._done: bool = True self._submitted: bool = False self._leaderboard: Leaderboard | None = None self._hints: HintProvider | None = None self._overfitting_accum: float = 0.0 self._active_combos: list[str] = [] self._traps_triggered: list[str] = [] self._mitigated_traps: set[str] = set() # ------------------------------------------------------------------ # OpenEnv API # ------------------------------------------------------------------ def reset(self, task_id: str = "easy_churn") -> Observation: self._task = get_task(task_id) self._cv_score = self._task.base_cv_score self._test_score = self._task.base_test_score self._applied = [] self._history = [] self._step_count = 0 self._max_steps = 10 self._done = False self._submitted = False self._overfitting_accum = 0.0 self._active_combos = [] self._traps_triggered = [] self._mitigated_traps = set() self._leaderboard = Leaderboard(task_id, list(self._task.ghost_scores)) self._hints = HintProvider(list(self._task.hints)) return self._observation(message="Environment reset. Choose your first action.") def step(self, action: Action) -> StepResponse: if self._task is None or self._done: raise RuntimeError("Environment not active. Call reset() first.") self._step_count += 1 info: dict[str, Any] = {} prev_cv = self._cv_score tag = action.tag() # --- Validate --- valid, msg = self._validate_action(action) if not valid: reward = Reward(total=-0.05, breakdown=RewardBreakdown(redundancy_penalty=-0.05)) info["error"] = msg obs = self._observation(message=f"Invalid action: {msg}") done = self._check_done() return StepResponse(observation=obs, reward=reward, done=done, info=info) # --- Apply action (mutates scores) --- self._apply_action(action, info) # --- Check failure-mode traps --- traps_this_step = self._check_traps(action, info) # --- Check strategy combos --- newly_completed = self._check_combos() if newly_completed: info["combos_completed"] = newly_completed # --- Contextual relevance --- relevance = self._task.context_relevance.get(tag) is_submit = action.action_type == ActionType.SUBMIT reward = compute_reward( prev_cv=prev_cv, new_cv=self._cv_score, new_test=self._test_score, action_tag=tag, expected_strategies=self._task.expected_strategies, already_applied=self._applied[:-1] if tag in self._applied else self._applied, is_submit=is_submit, submitted_test_score=self._test_score if is_submit else None, newly_completed_combos=newly_completed, context_relevance=relevance, traps_triggered_this_step=traps_this_step, ) self._history.append(tag) done = self._check_done() obs = self._observation(message=info.get("message", "")) return StepResponse(observation=obs, reward=reward, done=done, info=info) def state(self) -> EnvState: assert self._task is not None assert self._leaderboard is not None assert self._hints is not None return EnvState( task_id=self._task.task_id, step_count=self._step_count, max_steps=self._task.max_steps, done=self._done, cv_score=round(self._cv_score, 4), test_score=round(self._test_score, 4), applied_strategies=list(self._applied), strategy_history=list(self._history), leaderboard_rank=self._leaderboard.agent_rank(self._test_score), leaderboard=self._leaderboard.full_board(self._test_score), submitted=self._submitted, hint_count=self._hints.hints_given, active_combos=list(self._active_combos), traps_triggered=list(self._traps_triggered), ) # ------------------------------------------------------------------ # Observation # ------------------------------------------------------------------ def _observation(self, message: str = "") -> Observation: assert self._task is not None and self._leaderboard is not None return Observation( dataset_metadata=self._task.dataset_metadata, applied_strategies=list(self._applied), current_cv_score=round(self._cv_score, 4), leaderboard_rank=self._leaderboard.agent_rank(self._test_score), step_count=self._step_count, max_steps=self._task.max_steps, done=self._done, message=message, ) # ------------------------------------------------------------------ # Validation (hierarchical categories) # ------------------------------------------------------------------ def _validate_action(self, action: Action) -> tuple[bool, str]: at = action.action_type p = action.parameters if at == ActionType.PSEUDO_LABEL: iters = p.get("iterations") if iters is None or not isinstance(iters, int) or iters < 1: return False, "pseudo_label requires 'iterations' (int >= 1)" return True, "" if at == ActionType.SUBMIT: if self._submitted: return False, "Already submitted." if not any(s.startswith("train_model:") for s in self._applied): return False, "Cannot submit without training a model first." return True, "" if at == ActionType.INSPECT_TOP_SOLUTION: return True, "" key = get_param_key(at.value) if key: value = p.get(key) allowed = get_allowed_values(at.value) if value is None or value not in allowed: return False, f"{at.value} requires '{key}' in {allowed}" cat_err = validate_category(at.value, p.get("category"), value) if cat_err: return False, cat_err return True, "" # ------------------------------------------------------------------ # Action application with causal logic # ------------------------------------------------------------------ def _apply_action(self, action: Action, info: dict[str, Any]) -> None: assert self._task is not None and self._hints is not None tag = action.tag() modifiers = self._task.score_modifiers overfit_risk = self._task.overfitting_risk props = self._task.dataset_properties if action.action_type == ActionType.INSPECT_TOP_SOLUTION: hint = self._hints.next_hint() info["hint"] = hint info["message"] = f"Hint: {hint}" if tag not in self._applied: self._applied.append(tag) return if action.action_type == ActionType.SUBMIT: self._submitted = True self._done = True info["message"] = ( f"Submitted! Final test score: {self._test_score:.4f} " f"CV score: {self._cv_score:.4f}" ) return is_repeat = tag in self._applied # --- Track mitigation actions --- self._update_mitigations(tag, props) if tag in modifiers: cv_delta = modifiers[tag]["cv"] test_delta = modifiers[tag]["test"] if is_repeat: cv_delta *= 0.1 test_delta *= 0.1 else: self._applied.append(tag) overfit_extra = overfit_risk.get(tag, 0.0) self._overfitting_accum += overfit_extra cv_delta += overfit_extra self._cv_score = min(1.0, self._cv_score + cv_delta) self._test_score = min(1.0, self._test_score + test_delta) rank = self._leaderboard.agent_rank(self._test_score) if self._leaderboard else "?" info["message"] = f"Applied {action.full_tag()}. CV: {self._cv_score:.4f}, Rank: {rank}" else: if tag not in self._applied: self._applied.append(tag) info["message"] = f"Applied {action.full_tag()} (no modifier for this task)." def _update_mitigations(self, tag: str, props: Any) -> None: """Track which failure modes the agent has pre-emptively mitigated.""" if tag == "detect_shift:adversarial_validation" and props.has_shift: self._mitigated_traps.add("ignoring_shift") if tag == "detect_shift:remove_identifiers" and props.has_shift: self._mitigated_traps.add("ignoring_shift") if tag == "clean_data:remove_leaky_features" and props.has_leakage: self._mitigated_traps.add("keeping_leaky_feature") if tag == "feature_engineering:sin_cos_encoding" and props.has_spatial_data: self._mitigated_traps.add("raw_heading_without_sincos") if tag.startswith("augmentation:") and props.has_images: self._mitigated_traps.add("no_augmentation_on_images") # ------------------------------------------------------------------ # Failure-mode trap detection # ------------------------------------------------------------------ def _check_traps(self, action: Action, info: dict[str, Any]) -> list[str]: assert self._task is not None tag = action.tag() props = self._task.dataset_properties triggered: list[str] = [] for fm in self._task.failure_modes: if fm.name in self._traps_triggered: continue if fm.name in self._mitigated_traps: continue if tag != fm.trigger_tag: continue prop_val = getattr(props, fm.condition_field, None) if prop_val == fm.condition_value: self._cv_score = min(1.0, self._cv_score + fm.cv_effect) self._test_score = max(0.0, self._test_score + fm.test_effect) self._traps_triggered.append(fm.name) triggered.append(fm.name) prev_msg = info.get("message", "") info["message"] = f"{prev_msg} TRAP: {fm.message}" info.setdefault("traps", []).append({ "name": fm.name, "message": fm.message, "cv_effect": fm.cv_effect, "test_effect": fm.test_effect, }) return triggered # ------------------------------------------------------------------ # Combo detection # ------------------------------------------------------------------ def _check_combos(self) -> list[str]: assert self._task is not None newly_completed: list[str] = [] applied_set = set(self._applied) for combo in self._task.strategy_combos: if combo.name in self._active_combos: continue if combo.required.issubset(applied_set): self._active_combos.append(combo.name) self._cv_score = min(1.0, self._cv_score + combo.cv_bonus) self._test_score = min(1.0, self._test_score + combo.test_bonus) newly_completed.append(combo.name) return newly_completed # ------------------------------------------------------------------ # Done # ------------------------------------------------------------------ def _check_done(self) -> bool: if self._done: return True assert self._task is not None if self._step_count >= self._task.max_steps: self._done = True return self._done