rewards/engine.py

"""
Multi-Signal Reward Engine — Computes composite rewards for the DevOps RL agent.

Each action receives a multi-component reward based on success, progress,
efficiency, safety, and other signals. Returns both the total reward
and a detailed breakdown for logging and analysis.
"""

from __future__ import annotations

import re
from typing import Dict, List, Tuple

from executor.docker_executor import ExecutionResult
from fingerprint.classifier import ErrorFingerprinter
from scenarios.registry import Scenario


class RewardEngine:
    """Computes multi-signal rewards for agent actions.

    Reward components:
        - success:          +10.0 when scenario success_condition is met
        - correct_command:   +3.0 when action matches a hint command
        - progress:          +1.0 when error log changes (shorter/different)
        - efficiency_bonus:  +2.0 when solved in ≤ len(hint_commands) steps
        - invalid_command:   -2.0 when command is not in the whitelist
        - dangerous_command: -10.0 when command matches blocklist
        - no_progress:       -1.0 when error log is identical to previous
        - timeout:           -5.0 when command times out
        - repeated_command:  -1.5 when same command issued twice in episode
        - step_cost:         -0.2 per step (encourages efficiency)

    Usage:
        engine = RewardEngine()
        total, breakdown = engine.compute_reward(
            action="pip install flask",
            result=execution_result,
            scenario=scenario,
            step_count=1,
            command_history=["pip install flask"],
            prev_error_log="ModuleNotFoundError...",
            curr_error_log="Successfully installed flask",
        )
    """

    # Reward signal values (configurable)
    REWARD_SUCCESS: float = 10.0
    REWARD_CORRECT_COMMAND: float = 1.5
    REWARD_PROGRESS: float = 1.0
    REWARD_EFFICIENCY_BONUS: float = 2.0
    PENALTY_INVALID_COMMAND: float = -2.0
    PENALTY_DANGEROUS_COMMAND: float = -10.0
    PENALTY_NO_PROGRESS: float = -1.0
    PENALTY_TIMEOUT: float = -5.0
    PENALTY_REPEATED_COMMAND: float = -1.5
    PENALTY_STEP_COST: float = -0.2

    def __init__(self) -> None:
        """Initialize reward helpers."""
        self._fingerprinter = ErrorFingerprinter()

    def compute_reward(
        self,
        action: str,
        result: ExecutionResult,
        scenario: Scenario,
        step_count: int,
        command_history: List[str],
        prev_error_log: str,
        curr_error_log: str,
    ) -> Tuple[float, Dict[str, float]]:
        """Compute the multi-signal reward for an action.

        Args:
            action: The shell command that was executed.
            result: The execution result from the sandbox.
            scenario: The current scenario being solved.
            step_count: Current step number in the episode (1-indexed).
            command_history: All commands issued so far (including current).
            prev_error_log: Error log before this action.
            curr_error_log: Error log after this action.

        Returns:
            Tuple of (total_reward, breakdown_dict) where breakdown_dict
            maps signal names to their individual reward values.
        """
        breakdown: Dict[str, float] = {}
        action_stripped = action.strip()

        # 1. Step cost (always applied)
        breakdown["step_cost"] = self.PENALTY_STEP_COST

        # 2. Check for blocked/dangerous command
        if result.blocked:
            if "dangerous" in result.block_reason.lower() or "blocklist" in result.block_reason.lower():
                breakdown["dangerous_command"] = self.PENALTY_DANGEROUS_COMMAND
            else:
                breakdown["invalid_command"] = self.PENALTY_INVALID_COMMAND
            total = sum(breakdown.values())
            return total, breakdown

        # 3. Check for timeout
        if result.timed_out:
            breakdown["timeout"] = self.PENALTY_TIMEOUT
            total = sum(breakdown.values())
            return total, breakdown

        # 4. Check for repeated command
        if self._is_repeated(action_stripped, command_history):
            breakdown["repeated_command"] = self.PENALTY_REPEATED_COMMAND

        # 5. Check for progress
        made_progress = self._has_progress(prev_error_log, curr_error_log)
        if made_progress:
            breakdown["progress"] = self.REWARD_PROGRESS
        elif prev_error_log and curr_error_log and self._logs_identical(prev_error_log, curr_error_log):
            breakdown["no_progress"] = self.PENALTY_NO_PROGRESS

        # 6. Check for success
        combined_output = f"{result.stdout}\n{result.stderr}".strip()
        solved = scenario.success_condition(combined_output)
        if solved:
            breakdown["success"] = self.REWARD_SUCCESS

            # 7. Efficiency bonus
            if step_count <= len(scenario.hint_commands):
                breakdown["efficiency_bonus"] = self.REWARD_EFFICIENCY_BONUS

        # 8. Hint reward is only useful when accompanied by real improvement.
        if self._matches_hint(action_stripped, scenario.hint_commands) and (made_progress or solved):
            breakdown["correct_command"] = self.REWARD_CORRECT_COMMAND

        total = sum(breakdown.values())
        return total, breakdown

    def _is_repeated(self, action: str, command_history: List[str]) -> bool:
        """Check if the action was already issued in this episode.

        Args:
            action: Current action.
            command_history: All previous commands (not including current).

        Returns:
            True if the command was previously issued.
        """
        # command_history includes the current command, so check for >1 occurrence
        normalized = action.strip().lower()
        count = sum(1 for cmd in command_history if cmd.strip().lower() == normalized)
        return count > 1

    def _matches_hint(self, action: str, hint_commands: List[str]) -> bool:
        """Check if the action matches any hint command.

        Uses flexible matching: strips whitespace, normalizes separators,
        and checks for substring containment.

        Args:
            action: The command to check.
            hint_commands: List of optimal commands from the scenario.

        Returns:
            True if the action matches a hint command.
        """
        action_normalized = self._normalize_command(action)
        for hint in hint_commands:
            hint_normalized = self._normalize_command(hint)
            if action_normalized == hint_normalized:
                return True
            # Check if the core command is present (e.g., "pip install flask" in
            # "pip install flask==2.0")
            if hint_normalized in action_normalized or action_normalized in hint_normalized:
                return True
        return False

    def _normalize_command(self, cmd: str) -> str:
        """Normalize a command for comparison.

        Args:
            cmd: Command string to normalize.

        Returns:
            Normalized command string.
        """
        # Strip, lowercase, collapse whitespace
        normalized = cmd.strip().lower()
        normalized = re.sub(r'\s+', ' ', normalized)
        return normalized

    def _has_progress(self, prev_log: str, curr_log: str) -> bool:
        """Check if there has been progress (error changed or reduced).

        Args:
            prev_log: Previous error log.
            curr_log: Current error log.

        Returns:
            True if progress was made (error changed for the better).
        """
        if not prev_log:
            return False
        if not curr_log:
            return True  # Error cleared entirely

        prev_stripped = prev_log.strip()
        curr_stripped = curr_log.strip()
        curr_lower = curr_stripped.lower()

        if prev_stripped == curr_stripped:
            return False

        success_keywords = ["success", "installed", "running", "ok", "complete"]
        failure_keywords = ["traceback", "error", "exception", "failed", "not found", "cannot"]

        if any(kw in curr_lower for kw in success_keywords) and not any(kw in curr_lower for kw in failure_keywords):
            return True

        prev_fp = self._fingerprinter.classify(prev_stripped)
        curr_fp = self._fingerprinter.classify(curr_stripped)

        # Severity reduction: fewer hard-failure tokens means better state.
        if self._error_severity(curr_stripped) < self._error_severity(prev_stripped):
            return True

        # If the same error family remains, lower classifier confidence can indicate a weaker/fading failure signature.
        if prev_fp.error_type == curr_fp.error_type and curr_fp.confidence < prev_fp.confidence:
            return True

        # Reduced output while staying in the same error family can indicate partial remediation.
        if prev_fp.error_type == curr_fp.error_type and len(curr_stripped) < len(prev_stripped):
            return True

        # Resolved from known error to unknown/no-error-like output.
        if prev_fp.error_type != "unknown" and curr_fp.error_type == "unknown":
            if not any(kw in curr_lower for kw in failure_keywords):
                return True

        return False

    def _error_severity(self, log: str) -> int:
        """Estimate error severity from high-signal failure markers."""
        lowered = log.lower()
        markers = ["traceback", "exception", "error", "failed", "fatal", "cannot", "not found"]
        return sum(lowered.count(marker) for marker in markers)

    def _logs_identical(self, prev_log: str, curr_log: str) -> bool:
        """Check if two error logs are essentially identical.

        Args:
            prev_log: Previous error log.
            curr_log: Current error log.

        Returns:
            True if the logs are identical after normalization.
        """
        return prev_log.strip() == curr_log.strip()