env/reward.py

"""
Reward function - provides dense, shaped signals to guide learning.

Signal summary
--------------
+150   full resolution bonus
+ 30   correct intermediate fix step
- 15   wrong action (no progress)
-0.04  per ms of latency  (continuous cost)
- 25   per unit of error_rate (continuous cost)
-  2   time penalty per step (urgency)
"""
from __future__ import annotations
from typing import Dict, Any, Tuple


def compute_reward(
    prev: Dict[str, Any],
    curr: Dict[str, Any],
    action: str,
    time_step: int = 0,
) -> Tuple[float, Dict[str, Any]]:
    reward = 0.0

    # Resolution bonus
    if curr["resolved"]:
        reward += 150.0

    # Partial progress
    if curr["fix_progress"] > prev["fix_progress"]:
        reward += 30.0
    elif action != "noop":
        # Wrong action (no progress, not a passive noop)
        reward -= 15.0

    # Continuous metric penalties
    reward -= curr["metrics"]["latency"]    * 0.04
    reward -= curr["metrics"]["error_rate"] * 25.0

    # Time penalty (escalates after step 10 for urgency)
    time_penalty = 2.0 + (0.5 * max(0, time_step - 10))
    reward -= time_penalty

    info = {
        "latency":    round(curr["metrics"]["latency"], 2),
        "error_rate": round(curr["metrics"]["error_rate"], 4),
        "cpu":        round(curr["metrics"]["cpu"], 2),
        "progress":   curr["fix_progress"],
        "resolved":   curr["resolved"],
    }

    return reward, info