Spaces:
Sleeping
Sleeping
| """Reward architecture β multi-signal verifiable reward, no LLM judge.""" | |
| from __future__ import annotations | |
| import math | |
| import re | |
| from collections import Counter | |
| from typing import Any, Dict, List, Optional | |
| try: | |
| from models import ( | |
| FlakeForgeAction, | |
| FlakeForgeObservation, | |
| RewardBreakdown, | |
| ROOT_CAUSE_TYPES, | |
| RELATED_CATEGORIES, | |
| ThinkClaim, | |
| failure_mode_entropy, | |
| ) | |
| except ImportError: | |
| from ..models import ( | |
| FlakeForgeAction, | |
| FlakeForgeObservation, | |
| RewardBreakdown, | |
| ROOT_CAUSE_TYPES, | |
| RELATED_CATEGORIES, | |
| ThinkClaim, | |
| failure_mode_entropy, | |
| ) | |
| try: | |
| from agent.unified_agent import infer_category_from_patch | |
| except ImportError: | |
| try: | |
| from ..agent.unified_agent import infer_category_from_patch | |
| except ImportError: | |
| def infer_category_from_patch(patch_text: str) -> str: | |
| return "unknown" | |
| def _claim_has_required_format(claim: Any) -> bool: | |
| return ( | |
| getattr(claim, "category", "") in ROOT_CAUSE_TYPES | |
| and bool(str(getattr(claim, "entity", "")).strip()) | |
| and bool(str(getattr(claim, "location", "")).strip()) | |
| and getattr(claim, "polarity", "") in ("present", "absent") | |
| and bool(str(getattr(claim, "reason", "")).strip()) | |
| ) | |
| def _hunk_has_required_format(hunk: Any) -> bool: | |
| return ( | |
| bool(str(getattr(hunk, "file", "")).strip()) | |
| and bool(str(getattr(hunk, "search", "")).strip()) | |
| and isinstance(getattr(hunk, "replace", None), str) | |
| ) | |
| def compute_format_reward(action: FlakeForgeAction) -> float: | |
| """Score current JSON-first action format. Returns 0..1. | |
| The model-facing schema now requires only: | |
| - claim: category, entity, location, polarity, reason | |
| - hunk: file, search, replace | |
| Optional metadata such as claim_id, ast_node_type, hunk_id, rationale, and | |
| addresses_claim should not affect format reward. | |
| """ | |
| score = 0.0 | |
| think_score = 0.0 | |
| st = action.structured_think | |
| if st is not None: | |
| valid_claims = [claim for claim in st.claims if _claim_has_required_format(claim)] | |
| if valid_claims and st.format_penalty == 0.0: | |
| think_score = 0.5 | |
| elif valid_claims and st.format_penalty > -0.5: | |
| think_score = 0.4 | |
| elif st.claims: | |
| think_score = 0.25 | |
| if think_score == 0.0 and action.think_text: | |
| think = action.think_text.lower() | |
| if ("root cause:" in think or '"category"' in think) and ( | |
| "confidence" in think or '"reason"' in think | |
| ): | |
| think_score = 0.5 | |
| score += think_score | |
| patch_score = 0.0 | |
| sp = action.structured_patch | |
| if sp is not None: | |
| valid_hunks = [hunk for hunk in sp.hunks if _hunk_has_required_format(hunk)] | |
| if valid_hunks and sp.format_penalty == 0.0: | |
| patch_score = 0.5 | |
| elif valid_hunks and sp.format_penalty > -0.5: | |
| patch_score = 0.4 | |
| elif sp.hunks: | |
| patch_score = 0.25 | |
| if patch_score == 0.0 and action.patch_text: | |
| has_search = "<<<<<<< SEARCH" in action.patch_text or "<<<<<<<" in action.patch_text | |
| has_replace = ">>>>>>> REPLACE" in action.patch_text or ">>>>>>>" in action.patch_text | |
| has_sep = "=======" in action.patch_text | |
| if has_search and has_replace and has_sep: | |
| patch_score = 0.5 | |
| elif has_search or has_replace: | |
| patch_score = 0.25 | |
| score += patch_score | |
| return round(min(score, 1.0), 2) | |
| def compute_compile_reward( | |
| patch_applied_successfully: bool, | |
| syntax_error: Optional[str] = None, | |
| *, | |
| rejected_by_validator: bool = False, | |
| rolled_back: bool = False, | |
| ) -> float: | |
| """Score patch applicability and syntax correctness. Returns -1.0..1.0.""" | |
| if rejected_by_validator: | |
| return -0.5 | |
| if rolled_back: | |
| return 0.0 | |
| if not patch_applied_successfully: | |
| return -1.0 | |
| if syntax_error: | |
| return -0.5 | |
| return 1.0 | |
| def _potential(pass_rate: float) -> float: | |
| """Potential function Ξ¦(pass_rate) for potential-based reward shaping. | |
| Uses Ξ¦ = pass_rate^2 so improvements near 1.0 are worth more than | |
| improvements near 0.0. Theory (Ng, Harada, Russell 1999) guarantees | |
| that Ξ¦(s') - Ξ¦(s) preserves the optimal policy. | |
| """ | |
| return pass_rate ** 2 | |
| def compute_stability_reward( | |
| baseline_pass_rate: float, | |
| current_pass_rate: float, | |
| ) -> float: | |
| """Potential-based stability reward: Ξ¦(current) - Ξ¦(baseline). Returns ~-1.0..3.0.""" | |
| shaped = _potential(current_pass_rate) - _potential(baseline_pass_rate) | |
| if current_pass_rate >= 1.0: | |
| return max(2.0, round(shaped * 3.0, 4)) | |
| elif shaped > 0: | |
| return round(min(shaped * 3.0, 1.5), 4) | |
| elif shaped < 0: | |
| return round(max(shaped * 5.0, -1.0), 4) | |
| else: | |
| return -0.1 | |
| def compute_causal_proximity_reward( | |
| patch_files: List[str], | |
| failure_frontier: str, | |
| call_chain: List[str], | |
| boundary_crossings: List[str], | |
| ) -> float: | |
| """Reward patches that touch code near the failure site. Returns -0.3..1.0.""" | |
| if not patch_files or not failure_frontier: | |
| return 0.0 | |
| score = 0.0 | |
| frontier_file = failure_frontier.split(":")[0] if ":" in failure_frontier else failure_frontier | |
| for pf in patch_files: | |
| pf_name = pf.split("/")[-1] if "/" in pf else pf | |
| frontier_name = frontier_file.split("/")[-1] if "/" in frontier_file else frontier_file | |
| if pf_name == frontier_name: | |
| score = 1.0 | |
| break | |
| if score == 0.0: | |
| for pf in patch_files: | |
| pf_name = pf.split("/")[-1].replace(".py", "") | |
| for frame in call_chain: | |
| if pf_name in frame: | |
| score = max(score, 0.5) | |
| break | |
| if score == 0.0: | |
| score = -0.3 | |
| if boundary_crossings and score > 0: | |
| score = min(score + 0.2, 1.0) | |
| return round(score, 2) | |
| def compute_entropy_reward( | |
| pre_entropy: float, | |
| post_entropy: float, | |
| ) -> float: | |
| """Reward reduction in failure mode entropy. Returns -0.5..1.0.""" | |
| if pre_entropy <= 0 and post_entropy <= 0: | |
| return 0.5 | |
| delta = pre_entropy - post_entropy | |
| if delta > 0: | |
| return round(min(delta * 2.0, 1.0), 4) | |
| elif delta < 0: | |
| return round(max(delta * 1.5, -0.5), 4) | |
| return 0.0 | |
| def compute_anti_hack_penalty( | |
| patch_text: str, | |
| files_modified: List[str], | |
| lines_changed: int, | |
| ) -> float: | |
| """Penalise common reward-hacking patterns. Returns -2.0..0.0.""" | |
| penalty = 0.0 | |
| patch_lower = patch_text.lower() | |
| if re.search(r"\bassert\b", patch_lower): | |
| search_blocks = re.findall(r"<<<<<<< SEARCH\n(.*?)=======", patch_text, re.DOTALL) | |
| replace_blocks = re.findall(r"=======\n(.*?)>>>>>>> REPLACE", patch_text, re.DOTALL) | |
| for search, replace in zip(search_blocks, replace_blocks): | |
| search_asserts = len(re.findall(r"\bassert\b", search.lower())) | |
| replace_asserts = len(re.findall(r"\bassert\b", replace.lower())) | |
| if search_asserts > replace_asserts: | |
| penalty -= 0.5 * (search_asserts - replace_asserts) | |
| sleep_patterns = ["time.sleep(", "await asyncio.sleep(", "sleep("] | |
| for pattern in sleep_patterns: | |
| if pattern in patch_lower: | |
| replace_count = sum( | |
| 1 for b in re.findall(r"=======\n(.*?)>>>>>>> REPLACE", patch_text, re.DOTALL) | |
| if pattern in b.lower() | |
| ) | |
| search_count = sum( | |
| 1 for b in re.findall(r"<<<<<<< SEARCH\n(.*?)=======", patch_text, re.DOTALL) | |
| if pattern in b.lower() | |
| ) | |
| if replace_count > search_count: | |
| penalty -= 0.3 * (replace_count - search_count) | |
| if "except:" in patch_lower or "except Exception:" in patch_lower: | |
| replace_blocks = re.findall(r"=======\n(.*?)>>>>>>> REPLACE", patch_text, re.DOTALL) | |
| for block in replace_blocks: | |
| if "except:" in block.lower() or "except exception:" in block.lower(): | |
| penalty -= 0.5 | |
| skip_patterns = ["@pytest.mark.skip", "@unittest.skip", "pytest.skip("] | |
| for pattern in skip_patterns: | |
| if pattern in patch_lower: | |
| penalty -= 1.0 | |
| if lines_changed > 50: | |
| penalty -= min((lines_changed - 50) / 100.0, 0.5) | |
| return round(max(penalty, -2.0), 2) | |
| def compute_reasoning_consistency( | |
| predicted_category: str, | |
| inferred_category_from_patch: str, | |
| think_text: str, | |
| patch_text: str, | |
| ) -> float: | |
| """Category match between predicted and patch-inferred root cause. Returns -0.5..0.5.""" | |
| if not predicted_category or not inferred_category_from_patch: | |
| return 0.0 | |
| if predicted_category == inferred_category_from_patch: | |
| return 0.5 | |
| related = RELATED_CATEGORIES.get(predicted_category, set()) | |
| if inferred_category_from_patch in related: | |
| return 0.25 | |
| return -0.5 | |
| def compute_think_history_penalty( | |
| action: FlakeForgeAction, | |
| think_history: Optional[List[Dict[str, Any]]], | |
| ) -> float: | |
| """Mildly penalize repeating the same diagnosis without new evidence.""" | |
| if not think_history: | |
| return 0.0 | |
| categories: set[str] = set() | |
| entities: set[str] = set() | |
| reasons: set[str] = set() | |
| if action.structured_think and action.structured_think.claims: | |
| for claim in action.structured_think.claims: | |
| if claim.category: | |
| categories.add(claim.category) | |
| if claim.entity: | |
| entities.add(claim.entity) | |
| if claim.reason: | |
| reasons.add(claim.reason[:35].lower().strip()) | |
| elif action.predicted_category: | |
| categories.add(action.predicted_category) | |
| if not categories and not entities and not reasons: | |
| return 0.0 | |
| previous = think_history[-1] | |
| repeated_category = bool(categories & set(previous.get("categories", []))) | |
| repeated_entity = bool(entities & set(previous.get("entities", []))) | |
| repeated_reason = bool(reasons & set(previous.get("reason_signatures", []))) | |
| if repeated_category and (repeated_entity or repeated_reason): | |
| return -0.25 | |
| return 0.0 | |
| def compute_verifiable_reward( | |
| action: FlakeForgeAction, | |
| observation: FlakeForgeObservation, | |
| patch_result: Dict[str, Any], | |
| post_run_results: List[Dict[str, Any]], | |
| baseline_pass_rate: float, | |
| pre_entropy: float, | |
| oracle_score: Optional[float] = None, | |
| regression_detected: bool = False, | |
| think_history: Optional[List[Dict[str, Any]]] = None, | |
| ) -> RewardBreakdown: | |
| """Compute the full multi-signal verifiable reward.""" | |
| patch_applied = patch_result.get("success", False) | |
| rejected_by_validator = bool(patch_result.get("rejected_by_validator", False)) | |
| rolled_back = bool(patch_result.get("rolled_back", False)) | |
| syntax_error = patch_result.get("syntax_error") if patch_applied else None | |
| files_modified = patch_result.get("files_modified", []) | |
| lines_changed = patch_result.get("lines_changed", 0) | |
| noop_patch = bool(patch_result.get("noop", False)) | |
| regression_detected = bool(regression_detected or patch_result.get("regression_detected", False)) | |
| post_pass_count = sum(1 for r in post_run_results if r.get("passed", False)) | |
| if post_run_results: | |
| post_pass_rate = post_pass_count / len(post_run_results) | |
| post_errors = [r.get("error_type", "") for r in post_run_results if not r.get("passed", True)] | |
| post_counter = Counter(e for e in post_errors if e) | |
| if post_counter: | |
| total = sum(post_counter.values()) | |
| post_entropy = -sum((c / total) * math.log2(c / total) for c in post_counter.values()) | |
| max_e = math.log2(len(post_counter)) if len(post_counter) > 1 else 1.0 | |
| post_entropy = post_entropy / max_e if max_e > 0 else 0.0 | |
| else: | |
| post_entropy = 0.0 | |
| else: | |
| # If no tests were run, assume pass rate and entropy haven't changed | |
| post_pass_rate = observation.current_pass_rate | |
| post_entropy = pre_entropy | |
| inferred_cat = infer_category_from_patch(action.patch_text) | |
| breakdown = RewardBreakdown() | |
| # ββ Hard Gates Evaluation βββββββββββββββββββββββββββββββββββββββββββββββββ | |
| breakdown.format_reward = compute_format_reward(action) | |
| breakdown.compile_reward = compute_compile_reward( | |
| patch_applied, | |
| syntax_error, | |
| rejected_by_validator=rejected_by_validator, | |
| rolled_back=rolled_back, | |
| ) | |
| breakdown.regression_penalty = -1.5 if regression_detected else 0.0 | |
| breakdown.anti_hack_penalty = compute_anti_hack_penalty( | |
| action.patch_text, files_modified, lines_changed, | |
| ) | |
| # Fast-fail only for truly catastrophic issues (anti-hack or broken format). | |
| # Compile/validation failures still get partial credit so GRPO groups have | |
| # reward variance and can produce non-zero advantages. | |
| if breakdown.format_reward < 0.75 or breakdown.anti_hack_penalty < 0.0: | |
| breakdown.total_reward = -2.0 if breakdown.anti_hack_penalty < 0.0 else -1.0 | |
| return breakdown | |
| if breakdown.compile_reward < 1.0: | |
| breakdown.total_reward = round( | |
| breakdown.format_reward * 0.5 | |
| + breakdown.compile_reward * 1.0, | |
| 4, | |
| ) | |
| return breakdown | |
| # PatchValidator positive shaping | |
| if patch_applied and patch_result.get("validation_score") is not None: | |
| try: | |
| vs = float(patch_result["validation_score"]) | |
| breakdown.patch_validation_signal = round(0.2 * max(0.0, min(1.0, vs)), 4) | |
| except (TypeError, ValueError): | |
| breakdown.patch_validation_signal = 0.0 | |
| else: | |
| breakdown.patch_validation_signal = 0.0 | |
| breakdown.stability_reward = compute_stability_reward(baseline_pass_rate, post_pass_rate) | |
| breakdown.causal_proximity_reward = compute_causal_proximity_reward( | |
| files_modified, | |
| observation.failure_frontier, | |
| observation.call_chain_to_frontier, | |
| observation.boundary_crossings, | |
| ) | |
| breakdown.failure_entropy_reward = compute_entropy_reward(pre_entropy, post_entropy) | |
| breakdown.noop_patch_penalty = -0.5 if patch_applied and noop_patch else 0.0 | |
| breakdown.think_history_penalty = compute_think_history_penalty(action, think_history) | |
| # ββ Oracle (mild reasoning reward) ββββββββββββββ | |
| if oracle_score is not None: | |
| breakdown.oracle_reasoning_reward = round(float(oracle_score), 4) | |
| breakdown.reasoning_consistency_reward = 0.0 | |
| else: | |
| breakdown.reasoning_consistency_reward = compute_reasoning_consistency( | |
| action.predicted_category, inferred_cat, action.think_text, action.patch_text, | |
| ) | |
| breakdown.oracle_reasoning_reward = 0.0 | |
| # ββ Terminal bonus (scaled down to 4.0) βββββββββββββββββββββββββββββββββββ | |
| if patch_applied and post_run_results: | |
| if post_pass_rate >= 1.0: | |
| breakdown.terminal_bonus = 4.0 # Solving it = massive positive signal | |
| elif post_pass_rate > baseline_pass_rate + 0.5: | |
| breakdown.terminal_bonus = 2.0 # Major improvement | |
| elif post_pass_rate > baseline_pass_rate + 0.3: | |
| breakdown.terminal_bonus = 1.0 # Meaningful progress | |
| # ββ Total reward ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ | |
| oracle_component = ( | |
| breakdown.oracle_reasoning_reward * 1.0 # reduced weight 1.0 (was 2.5) | |
| if oracle_score is not None | |
| else breakdown.reasoning_consistency_reward * 0.5 | |
| ) | |
| breakdown.total_reward = round( | |
| breakdown.format_reward * 0.5 | |
| + breakdown.compile_reward * 1.0 | |
| + breakdown.stability_reward * 3.0 # increased weight 3.0 (was 2.0) | |
| + breakdown.causal_proximity_reward * 0.5 | |
| + breakdown.failure_entropy_reward * 0.5 | |
| + breakdown.anti_hack_penalty * 1.5 | |
| + breakdown.regression_penalty * 1.0 | |
| + oracle_component | |
| + breakdown.patch_validation_signal * 1.0 | |
| + breakdown.noop_patch_penalty * 1.0 | |
| + breakdown.think_history_penalty * 1.0 | |
| + breakdown.terminal_bonus * 1.0, | |
| 4, | |
| ) | |
| return breakdown | |