ahagestedt/apex-devops
0
1"""2Shaped reward system for QA evaluation environments.3 4Implements a 7-component reward function that incentivizes agents to:51. Gather complete information before making decisions (completeness)62. Score accurately against ground truth (accuracy)73. Provide high-quality reasoning (reasoning_quality)84. Calibrate confidence properly (confidence_calibration)95. Catch safety violations (safety_compliance)106. Be efficient with fewer steps (efficiency)117. Route correctly via auto_commit vs route_to_hitl (gating_correctness)12 13Total reward range: [-0.3, 1.0] (matching clinKriya reference environment).14"""15 16from __future__ import annotations17 18from typing import Any, Dict, List, Optional19 20 21# ---------------------------------------------------------------------------22# Constants23# ---------------------------------------------------------------------------24 25AVAILABLE_RESOURCES: tuple[str, ...] = (26 "ticket",27 "customer_history",28 "rubric",29 "similar_cases",30 "api_docs",31)32 33COMPONENT_WEIGHTS: Dict[str, float] = {34 "completeness": 0.20,35 "accuracy": 0.25,36 "reasoning_quality": 0.15,37 "confidence_calibration": 0.10,38 "safety_compliance": 0.15,39 "efficiency": 0.05,40 "gating_correctness": 0.10,41}42 43COMPONENT_RANGES: Dict[str, tuple[float, float]] = {44 "completeness": (0.0, 1.0),45 "accuracy": (0.0, 1.0),46 "reasoning_quality": (0.0, 1.0),47 "confidence_calibration": (0.0, 1.0),48 "safety_compliance": (-1.0, 1.0),49 "efficiency": (0.0, 1.0),50 "gating_correctness": (0.0, 1.0),51}52 53TOTAL_REWARD_MIN: float = -0.354TOTAL_REWARD_MAX: float = 1.055 56 57def _clamp(value: float, low: float, high: float) -> float:58 """Clamp a value to [low, high]."""59 return max(low, min(high, value))60 61 62# ---------------------------------------------------------------------------63# RewardCalculator64# ---------------------------------------------------------------------------65 66 67class RewardCalculator:68 """Compute a 7-component shaped reward for QA evaluation episodes.69 70 Each component is computed independently, multiplied by its weight,71 and summed to produce a total reward clamped to [-0.3, 1.0].72 73 Attributes:74 AUTO_COMMIT_THRESHOLD: Minimum confidence for auto-commit gating.75 MAX_STEPS: Maximum expected steps in an episode.76 """77 78 AUTO_COMMIT_THRESHOLD: float = 0.8579 MAX_STEPS: int = 880 81 def compute_reward(self, state_data: dict[str, Any]) -> dict[str, Any]:82 """Compute all 7 reward components and the total shaped reward.83 84 Args:85 state_data: Episode data with keys:86 - gathered_info: dict of resources the agent retrieved87 (keys are resource names, e.g. "ticket", "rubric").88 - submitted_evaluation: dict with the agent's score (int),89 confidence (float), reasoning (str), and optionally90 safety_flag (bool) and reasoning_steps (list).91 - ground_truth: dict with score (int) and92 safety_violation (bool).93 - action_history: list of action dicts taken during the94 episode.95 - scenario: full scenario definition (unused currently,96 reserved for future scenario-aware shaping).97 98 Returns:99 dict with:100 - components: raw (unweighted) value per component.101 - weighted_components: value * weight per component.102 - total: final reward clamped to [-0.3, 1.0].103 - gating_decision: "auto_commit" or "route_to_hitl".104 """105 gathered_info: dict[str, Any] = state_data.get("gathered_info", {})106 submitted: dict[str, Any] = state_data.get("submitted_evaluation", {})107 ground_truth: dict[str, Any] = state_data.get("ground_truth", {})108 action_history: list[dict[str, Any]] = state_data.get("action_history", [])109 110 # Derive values with safe defaults111 submitted_score: int = int(submitted.get("score", 0))112 ground_truth_score: int = int(ground_truth.get("score", 0))113 submitted_confidence: float = float(submitted.get("confidence", 0.0))114 steps_taken: int = len(action_history)115 116 # Determine the phase index: how many actions occurred before the117 # first "submit" or "evaluate" action. This tells us whether the118 # agent gathered info *before* deciding.119 phase_index: int = self._find_decision_phase_index(action_history)120 121 # --- Compute each component ---122 accuracy_score = self._compute_accuracy(submitted_score, ground_truth_score)123 124 components: Dict[str, float] = {125 "completeness": self._compute_completeness(gathered_info, phase_index),126 "accuracy": accuracy_score,127 "reasoning_quality": self._compute_reasoning_quality(submitted),128 "confidence_calibration": self._compute_confidence_calibration(129 submitted_confidence, accuracy_score130 ),131 "safety_compliance": self._compute_safety_compliance(submitted, ground_truth),132 "efficiency": self._compute_efficiency(steps_taken),133 "gating_correctness": self._compute_gating_correctness(134 submitted_confidence, ground_truth_score135 ),136 }137 138 # Clamp each component to its valid range139 for name, value in components.items():140 low, high = COMPONENT_RANGES[name]141 components[name] = _clamp(value, low, high)142 143 # Apply weights144 weighted: Dict[str, float] = {145 name: round(components[name] * COMPONENT_WEIGHTS[name], 6)146 for name in components147 }148 149 total = _clamp(sum(weighted.values()), TOTAL_REWARD_MIN, TOTAL_REWARD_MAX)150 151 gating_decision = self.compute_gating(submitted_confidence)152 153 return {154 "components": components,155 "weighted_components": weighted,156 "total": round(total, 4),157 "gating_decision": gating_decision,158 }159 160 # ------------------------------------------------------------------161 # Individual component methods162 # ------------------------------------------------------------------163 164 def _compute_completeness(165 self, gathered_info: dict[str, Any], phase_index: int166 ) -> float:167 """Score based on how much information the agent gathered.168 169 Available resources: ticket, customer_history, rubric,170 similar_cases, api_docs.171 172 Base score = resources_gathered / total_available.173 A 0.1 bonus is added if all gathering happened before the174 decision action (phase_index indicates the agent looked at175 everything before committing).176 177 Args:178 gathered_info: Mapping of resource_name -> content/data179 retrieved by the agent.180 phase_index: Number of actions before the first decision181 action. A higher value relative to gathered count182 indicates gathering happened first.183 184 Returns:185 Float in [0, 1].186 """187 if not AVAILABLE_RESOURCES:188 return 1.0189 190 total_available = len(AVAILABLE_RESOURCES)191 gathered_count = sum(192 1 for resource in AVAILABLE_RESOURCES if resource in gathered_info193 )194 195 base_score = gathered_count / total_available196 197 # Bonus: if the agent gathered resources before making a decision.198 # phase_index >= gathered_count means all gathers preceded the199 # decision action.200 gather_before_decide_bonus = 0.1 if (201 gathered_count > 0 and phase_index >= gathered_count202 ) else 0.0203 204 return _clamp(base_score + gather_before_decide_bonus, 0.0, 1.0)205 206 def _compute_accuracy(207 self, submitted_score: int, ground_truth_score: int208 ) -> float:209 """Score accuracy: 1.0 for exact match, -0.25 per point of difference.210 211 Both scores are expected on a 1-5 integer scale. If the agent212 did not submit a score (0), this returns 0.0.213 214 Args:215 submitted_score: The agent's evaluation score.216 ground_truth_score: The correct score.217 218 Returns:219 Float in [0, 1].220 """221 if submitted_score == 0 or ground_truth_score == 0:222 return 0.0223 224 diff = abs(submitted_score - ground_truth_score)225 return _clamp(1.0 - 0.25 * diff, 0.0, 1.0)226 227 def _compute_reasoning_quality(228 self, submitted_evaluation: dict[str, Any]229 ) -> float:230 """Score the depth and coherence of the agent's reasoning.231 232 Criteria (additive):233 - Has non-empty reasoning text: +0.3234 - Reasoning length > 50 characters: +0.2235 - Has confidence value: +0.2236 - Has multiple reasoning steps (list with len >= 2): +0.3237 238 Args:239 submitted_evaluation: The agent's submitted evaluation dict,240 expected to contain "reasoning" (str), "confidence"241 (float), and optionally "reasoning_steps" (list).242 243 Returns:244 Float in [0, 1].245 """246 score = 0.0247 248 reasoning: str = str(submitted_evaluation.get("reasoning", ""))249 if reasoning.strip():250 score += 0.3251 252 if len(reasoning) > 50:253 score += 0.2254 255 confidence = submitted_evaluation.get("confidence")256 if confidence is not None:257 score += 0.2258 259 reasoning_steps: Optional[List[Any]] = submitted_evaluation.get(260 "reasoning_steps"261 )262 if reasoning_steps is not None and len(reasoning_steps) >= 2:263 score += 0.3264 265 return _clamp(score, 0.0, 1.0)266 267 def _compute_confidence_calibration(268 self, submitted_confidence: float, accuracy_score: float269 ) -> float:270 """Measure how well the agent's confidence matches its accuracy.271 272 Perfect calibration (confidence == accuracy_score) yields 1.0.273 The penalty is the absolute difference between the two.274 275 Args:276 submitted_confidence: The agent's stated confidence [0, 1].277 accuracy_score: The computed accuracy component [0, 1].278 279 Returns:280 Float in [0, 1].281 """282 return _clamp(1.0 - abs(submitted_confidence - accuracy_score), 0.0, 1.0)283 284 def _compute_safety_compliance(285 self, submitted_evaluation: dict[str, Any], ground_truth: dict[str, Any]286 ) -> float:287 """Reward or penalize safety violation detection.288 289 If ground_truth has safety_violation=True:290 - Agent flagged it (safety_flag=True): +1.0291 - Agent missed it: -1.0292 293 If ground_truth has safety_violation=False:294 - Agent correctly did not flag: +0.5295 - Agent falsely flagged: -0.3296 297 Args:298 submitted_evaluation: Agent's evaluation; should contain299 "safety_flag" (bool) if the agent flagged a violation.300 ground_truth: Must contain "safety_violation" (bool).301 302 Returns:303 Float in [-1, 1].304 """305 gt_violation: bool = bool(ground_truth.get("safety_violation", False))306 agent_flagged: bool = bool(submitted_evaluation.get("safety_flag", False))307 308 if gt_violation:309 return 1.0 if agent_flagged else -1.0310 else:311 return -0.3 if agent_flagged else 0.5312 313 def _compute_efficiency(self, steps_taken: int) -> float:314 """Score efficiency: fewer steps yield a higher score.315 316 Optimal is 3-4 steps. The formula treats 3 as the baseline:317 score = 1.0 - (steps_taken - 3) / (MAX_STEPS - 3)318 319 Steps at or below 3 get a perfect 1.0. Steps at MAX_STEPS320 get 0.0. Steps beyond MAX_STEPS are clamped to 0.0.321 322 Args:323 steps_taken: Number of actions in the episode.324 325 Returns:326 Float in [0, 1].327 """328 optimal_steps = 3329 denominator = self.MAX_STEPS - optimal_steps330 if denominator <= 0:331 return 1.0332 333 raw = 1.0 - (steps_taken - optimal_steps) / denominator334 return _clamp(raw, 0.0, 1.0)335 336 def _compute_gating_correctness(337 self, confidence: float, ground_truth_score: int338 ) -> float:339 """Score whether the agent's implicit gating matches the ideal.340 341 Ideal gating logic:342 auto_commit if confidence >= 0.85 AND score >= 4343 route_to_hitl otherwise344 345 The agent's actual gating decision is derived from its346 confidence. This component checks whether that decision347 matches what *should* happen given the ground truth score.348 349 Args:350 confidence: The agent's stated confidence [0, 1].351 ground_truth_score: The correct evaluation score (1-5).352 353 Returns:354 1.0 if the gating decision matches the ideal, 0.0 otherwise.355 """356 agent_decision = self.compute_gating(confidence)357 358 ideal_auto_commit = (359 confidence >= self.AUTO_COMMIT_THRESHOLD360 and ground_truth_score >= 4361 )362 ideal_decision = "auto_commit" if ideal_auto_commit else "route_to_hitl"363 364 return 1.0 if agent_decision == ideal_decision else 0.0365 366 def compute_gating(self, confidence: float) -> str:367 """Route based on confidence threshold.368 369 Args:370 confidence: The agent's confidence score [0, 1].371 372 Returns:373 "auto_commit" if confidence >= AUTO_COMMIT_THRESHOLD,374 "route_to_hitl" otherwise.375 """376 if confidence >= self.AUTO_COMMIT_THRESHOLD:377 return "auto_commit"378 return "route_to_hitl"379 380 # ------------------------------------------------------------------381 # Internal helpers382 # ------------------------------------------------------------------383 384 @staticmethod385 def _find_decision_phase_index(action_history: list[dict[str, Any]]) -> int:386 """Find the index of the first decision action in the history.387 388 A "decision" action is one whose action_type is "submit",389 "evaluate", or "finish". Everything before that index is390 considered the information-gathering phase.391 392 Args:393 action_history: List of action dicts, each expected to have394 an "action_type" key.395 396 Returns:397 The index of the first decision action, or len(action_history)398 if no decision action was found (i.e. all actions were399 gathering).400 """401 decision_types = {"submit", "evaluate", "finish"}402 for i, action in enumerate(action_history):403 if action.get("action_type", "") in decision_types:404 return i405 return len(action_history)406