from __future__ import annotations from typing import Any from env.anti_hacking import AntiHackingDetector from env.graders.deterministic import DeterministicGrader from env.hidden_tests import HiddenTestRunner class RewardCalculator: """Composes progress, execution, quality, and anti-hacking penalties.""" ACTION_PROGRESS_REWARDS = { "read_file": 0.02, "read_logs": 0.03, "analyze_error": 0.05, "edit_config": 0.06, "run_pipeline_stage": 0.07, "run_tests": 0.08, "validate_fix": 0.10, "submit_solution": 0.12, } QUALITY_WEIGHTS = { "deterministic": 0.40, "hidden": 0.25, "llm": 0.20, } def __init__( self, llm_judge: Any | None = None, anti_hacking_detector: AntiHackingDetector | None = None, deterministic_grader: DeterministicGrader | None = None, hidden_test_runner: HiddenTestRunner | None = None, ): self.llm_judge = llm_judge self.anti_hacking_detector = anti_hacking_detector or AntiHackingDetector() self.deterministic_grader = deterministic_grader or DeterministicGrader() self.hidden_test_runner = hidden_test_runner or HiddenTestRunner(grader=self.deterministic_grader) def calculate_step_reward( self, state: dict[str, Any] | None, action: str, result: dict[str, Any] | None, original_config: str | None = None, fixed_config: str | None = None, error_message: str | None = None, expected_config: str | None = None, metadata: dict[str, Any] | None = None, ) -> tuple[float, dict[str, float]]: state = state or {} result = result or {} metadata = metadata or {} current_config = fixed_config or result.get("fixed_config") or result.get("current_config") or "" expected_config = expected_config or result.get("expected_config") or state.get("expected_config") or "" original_config = original_config or result.get("original_config") or state.get("original_config") or "" error_message = error_message or result.get("error") or state.get("error") or "" prog = self._progress_reward(action, result) exec = self._execution_reward(result) det_score, hide_score, llm_score, qual = self._quality_reward( action=action, current_config=current_config, expected_config=expected_config, original_config=original_config, error_message=error_message, result=result, metadata=metadata, ) pen = self._penalty_reward(state=state, result=result, current_config=current_config) reward = prog + exec + qual + pen reward_clamped = round(self._clamp_01(reward), 4) components = { "progress": prog, "execution": exec, "deterministic": float(det_score or 0.0), "hidden": float(hide_score or 0.0), "llm_judge": float(llm_score or 0.0), "penalty": pen, "total": reward_clamped } return reward_clamped, components def _progress_reward(self, action: str, result: dict[str, Any]) -> float: reward = self.ACTION_PROGRESS_REWARDS.get(action, 0.0) if result.get("logs_analyzed"): reward += 0.04 if result.get("error_diagnosed"): reward += 0.08 if result.get("fix_proposed"): reward += 0.05 return reward def _execution_reward(self, result: dict[str, Any]) -> float: reward = 0.0 if result.get("pipeline_run"): reward += 0.10 if result.get("tests_passed"): reward += 0.20 if result.get("command_succeeded"): reward += 0.06 return reward def _quality_reward( self, action: str, current_config: str, expected_config: str, original_config: str, error_message: str, result: dict[str, Any], metadata: dict[str, Any], ) -> tuple[float, float, float, float]: if not current_config or not expected_config: return 0.0, 0.0, 0.0, 0.0010101 deterministic_score = result.get("deterministic_score") if deterministic_score is None: deterministic_score = self.deterministic_grader.grade(current_config, expected_config, metadata) hidden_pass_rate = result.get("hidden_test_pass_rate") if hidden_pass_rate is None and action in {"validate_fix", "submit_solution"}: hidden_pass_rate = self.hidden_test_runner.evaluate_fix( fixed_config=current_config, expected_config=expected_config, metadata=metadata, ) llm_average = 0.0 judge_scores = result.get("judge_scores") if not judge_scores and self.llm_judge and original_config and current_config: try: judge_scores = self.llm_judge.evaluate_fix(original_config, current_config, error_message) except Exception: judge_scores = None if isinstance(judge_scores, dict): correctness = self._clamp_01(judge_scores.get("correctness", 0.0)) minimalism = self._clamp_01(judge_scores.get("minimalism", 0.0)) quality = self._clamp_01(judge_scores.get("quality", 0.0)) llm_average = (correctness + minimalism + quality) / 3.0 quality_reward = 0.0 quality_reward += self.QUALITY_WEIGHTS["deterministic"] * self._clamp_01(deterministic_score) quality_reward += self.QUALITY_WEIGHTS["hidden"] * self._clamp_01(hidden_pass_rate or 0.0) quality_reward += self.QUALITY_WEIGHTS["llm"] * self._clamp_01(llm_average) return self._clamp_01(deterministic_score), self._clamp_01(hidden_pass_rate or 0.0), self._clamp_01(llm_average), quality_reward def _penalty_reward(self, state: dict[str, Any], result: dict[str, Any], current_config: str) -> float: changed_files_count = int(result.get("changed_files_count", state.get("changed_files_count", 0)) or 0) changed_lines_count = int(result.get("changed_lines_count", state.get("changed_lines_count", 0)) or 0) edit_count = result.get("edit_count", state.get("edit_count", 0)) step_count = int(state.get("step_count", 0) or 0) previous_config = result.get("previous_config") or state.get("previous_config") or "" consecutive_edit_actions = int( result.get("consecutive_edit_actions", state.get("consecutive_edit_actions", 0)) or 0 ) failed_validations = int(result.get("failed_validations", state.get("failed_validations", 0)) or 0) penalty = self.anti_hacking_detector.total_penalty( current_config=current_config, previous_config=previous_config, edit_count=edit_count, changed_files_count=changed_files_count, changed_lines_count=changed_lines_count, step_count=step_count, consecutive_edit_actions=consecutive_edit_actions, failed_validations=failed_validations, ) if result.get("hacking_attempt"): penalty -= 0.30 return penalty def _clamp_01(self, value: Any) -> float: try: parsed = float(value) except (TypeError, ValueError): parsed = 0.0 return max(0.001, min(0.999, float(parsed)))