server/rubrics.py

"""
Composite reward rubrics for Pyre (single-agent).

Each rubric class exposes a score() method.
The environment composes them by calling each rubric each step.

Per-step rubrics:
  TimeStepPenalty            -0.01      constant time pressure
  ProgressReward             +0.25      agent moved closer to nearest unblocked exit (BFS distance)
  ProgressRegressionPenalty  -0.15      agent moved farther from nearest exit (symmetric gradient)
  SafeProgressBonus          +0.05      stacks on ProgressReward when progress made through smoke-free cell
  DangerPenalty              -0.5       agent moved into smoke≥moderate or fire-adjacent cell
  HealthDrainPenalty         -0.02×dmg  proportional to health lost this step
  StrategicDoorBonus         +0.5       closed a door adjacent to active fire (once per door per episode)
  ExplorationBonus           +0.02      first visit to a cell in this episode

Episode-end rubrics:
  SelfSurviveBonus         +5.0              agent reached an open exit alive
  HealthSurvivalBonus      +1.5×(hp/100)     graduated bonus for evacuating with more health
  SelfDeathPenalty        -10.0              agent died (health depleted or fire/smoke)
  TimeoutPenalty           -5.0              agent ran out of steps while still alive (lighter than death)
  NearMissBonus            0→+3.0            partial credit on failure based on closest approach to exit
  TimeBonus                +0.05×rem         reward for finishing quickly (remaining steps)
"""

from collections import deque
from typing import Any, Dict, List, Optional, Set

from .fire_sim import EXIT_BLOCKED_FIRE_THRESHOLD, FIRE_BURNING

EXIT = 4
WALL = 1
OBSTACLE = 5
DOOR_CLOSED = 3
SMOKE_MODERATE = 0.5

BFS_INF = 9999
_BFS_INF = BFS_INF  # keep private alias for internal use


def unblocked_exits(exit_positions: List[List[int]], fire_grid: List[float], w: int) -> List[List[int]]:
    """Return exits that do not have significant fire on them."""
    return [
        ex for ex in exit_positions
        if fire_grid[ex[1] * w + ex[0]] < EXIT_BLOCKED_FIRE_THRESHOLD
    ]


# Legacy private alias so existing internal callers don't break.
_unblocked_exits = unblocked_exits


def bfs_exit_dist(
    x: int,
    y: int,
    exits: List[List[int]],
    cell_grid: List[int],
    w: int,
    h: int,
) -> int:
    """BFS traversal distance from (x, y) to the nearest reachable exit.

    Walls (1) and obstacles (5) block movement. Closed doors (3) are treated
    as passable — the agent can open them en route. Returns BFS_INF when no
    exit is reachable (all paths wall-blocked).
    """
    if not exits:
        return BFS_INF

    exit_set = {(ex[0], ex[1]) for ex in exits}
    if (x, y) in exit_set:
        return 0

    visited = {(x, y)}
    queue: deque = deque()
    queue.append((x, y, 0))

    while queue:
        cx, cy, dist = queue.popleft()
        for dx, dy in ((0, -1), (0, 1), (-1, 0), (1, 0)):
            nx, ny = cx + dx, cy + dy
            if not (0 <= nx < w and 0 <= ny < h):
                continue
            if (nx, ny) in visited:
                continue
            ct = cell_grid[ny * w + nx]
            if ct in (WALL, OBSTACLE):
                continue
            if (nx, ny) in exit_set:
                return dist + 1
            visited.add((nx, ny))
            queue.append((nx, ny, dist + 1))

    return BFS_INF


# Legacy private alias so existing internal callers don't break.
_bfs_exit_dist = bfs_exit_dist


# ---------------------------------------------------------------------------
# Per-step rubrics
# ---------------------------------------------------------------------------

class TimeStepPenalty:
    """Small constant penalty per step to encourage urgency."""

    def score(self, **_) -> float:
        return -0.01


class ProgressReward:
    """Reward agent for moving strictly closer to the nearest unblocked exit.

    Uses BFS traversal distance (respects walls and obstacles) instead of
    Manhattan distance, so only genuine navigational progress is rewarded.
    Value raised to +0.25 to create a stronger pull toward exits relative to
    the danger/loop penalties that push the agent away from threats.
    """

    def score(
        self,
        prev_agent_x: int, prev_agent_y: int,
        agent_x: int, agent_y: int,
        exit_positions: List[List[int]],
        fire_grid: List[float],
        cell_grid: List[int],
        w: int, h: int,
        action: str,
        **_,
    ) -> float:
        if action != "move":
            return 0.0
        exits = _unblocked_exits(exit_positions, fire_grid, w)
        if not exits:
            exits = exit_positions  # all blocked — still try to reward progress
        prev_dist = _bfs_exit_dist(prev_agent_x, prev_agent_y, exits, cell_grid, w, h)
        new_dist = _bfs_exit_dist(agent_x, agent_y, exits, cell_grid, w, h)
        return 0.25 if new_dist < prev_dist else 0.0


class DangerPenalty:
    """Penalise moving into a dangerous cell (smoke ≥ moderate or fire adjacent)."""

    def score(
        self,
        agent_x: int, agent_y: int,
        action: str,
        cell_grid: List[int],
        fire_grid: List[float],
        smoke_grid: List[float],
        w: int, h: int,
        **_,
    ) -> float:
        if action != "move":
            return 0.0

        i = agent_y * w + agent_x
        if i < 0 or i >= len(smoke_grid):
            return 0.0

        if smoke_grid[i] >= SMOKE_MODERATE:
            return -0.5

        for dx, dy in [(0, -1), (0, 1), (-1, 0), (1, 0)]:
            nx, ny = agent_x + dx, agent_y + dy
            if 0 <= nx < w and 0 <= ny < h:
                if fire_grid[ny * w + nx] >= FIRE_BURNING:
                    return -0.5
        return 0.0


class HealthDrainPenalty:
    """Penalty proportional to health damage taken this step from smoke/fire.

    Moderate smoke (~2 dmg/step) → -0.04/step
    Heavy smoke   (~5 dmg/step) → -0.10/step
    On fire       (~10 dmg/step) → -0.20/step
    """

    def score(self, health_damage: float, **_) -> float:
        return -0.02 * health_damage


class StrategicDoorBonus:
    """Bonus for closing a door adjacent to active fire — slows spread significantly.

    Each door can only earn this bonus once per episode. The environment passes
    a mutable `rewarded_doors` set; the rubric checks and updates it to prevent
    the agent from farming +0.5 by repeatedly opening and closing the same door.
    """

    def score(
        self,
        action: str,
        door_state: Optional[str],
        target_id: Optional[str],
        door_registry: Dict[str, List[int]],
        fire_grid: List[float],
        rewarded_doors: Set[str],
        w: int, h: int,
        **_,
    ) -> float:
        if action != "door" or door_state != "close" or not target_id:
            return 0.0
        if target_id not in door_registry:
            return 0.0
        if target_id in rewarded_doors:
            return 0.0

        dx, dy = door_registry[target_id]
        for ddx, ddy in [(0, -1), (0, 1), (-1, 0), (1, 0)]:
            nx, ny = dx + ddx, dy + ddy
            if 0 <= nx < w and 0 <= ny < h:
                if fire_grid[ny * w + nx] >= FIRE_BURNING:
                    rewarded_doors.add(target_id)
                    return 0.5
        return 0.0


class ProgressRegressionPenalty:
    """Penalise moving farther from the nearest unblocked exit.

    Symmetric counterpart to ProgressReward: the agent gets +0.25 for progress
    and –0.15 for regression, creating a strong two-sided gradient that
    discourages wandering away from the exit under fire pressure.
    """

    def score(
        self,
        prev_agent_x: int, prev_agent_y: int,
        agent_x: int, agent_y: int,
        exit_positions: List[List[int]],
        fire_grid: List[float],
        cell_grid: List[int],
        w: int, h: int,
        action: str,
        **_,
    ) -> float:
        if action != "move":
            return 0.0
        exits = unblocked_exits(exit_positions, fire_grid, w)
        if not exits:
            exits = exit_positions
        prev_dist = bfs_exit_dist(prev_agent_x, prev_agent_y, exits, cell_grid, w, h)
        new_dist = bfs_exit_dist(agent_x, agent_y, exits, cell_grid, w, h)
        return -0.15 if new_dist > prev_dist else 0.0


class SafeProgressBonus:
    """Extra reward for making exit-progress through a smoke-free cell.

    Stacks on top of ProgressReward to teach the agent to prefer safe routes
    when multiple paths lead equally close to the exit.
    """

    def score(
        self,
        prev_agent_x: int, prev_agent_y: int,
        agent_x: int, agent_y: int,
        exit_positions: List[List[int]],
        fire_grid: List[float],
        smoke_grid: List[float],
        cell_grid: List[int],
        w: int, h: int,
        action: str,
        **_,
    ) -> float:
        if action != "move":
            return 0.0
        exits = unblocked_exits(exit_positions, fire_grid, w)
        if not exits:
            exits = exit_positions
        prev_dist = bfs_exit_dist(prev_agent_x, prev_agent_y, exits, cell_grid, w, h)
        new_dist = bfs_exit_dist(agent_x, agent_y, exits, cell_grid, w, h)
        if new_dist >= prev_dist:
            return 0.0
        i = agent_y * w + agent_x
        return 0.05 if (0 <= i < len(smoke_grid) and smoke_grid[i] < SMOKE_MODERATE) else 0.0


class ExplorationBonus:
    """Small bonus for visiting a cell for the first time in the episode.

    Prevents the agent standing still or looping to avoid DangerPenalty.
    The environment tracks visited cells and passes `is_new_cell`.
    """

    def score(self, action: str, is_new_cell: bool, **_) -> float:
        return 0.02 if (action == "move" and is_new_cell) else 0.0


# ---------------------------------------------------------------------------
# Episode-end rubrics
# ---------------------------------------------------------------------------

class SelfSurviveBonus:
    """Big bonus when agent evacuated alive."""

    def score(self, done: bool, agent_evacuated: bool, **_) -> float:
        return 5.0 if (done and agent_evacuated) else 0.0


class HealthSurvivalBonus:
    """Graduated bonus for evacuating with remaining health.

    Rewards finding the *safest* route, not just any route. An agent that
    evacuates at 95 HP earns ~+1.43; one that barely survives at 5 HP earns
    ~+0.075. Range: 0 → +1.5.
    """

    def score(self, done: bool, agent_evacuated: bool, agent_health: float, **_) -> float:
        return 1.5 * (agent_health / 100.0) if (done and agent_evacuated) else 0.0


class SelfDeathPenalty:
    """Big penalty when agent died (health depleted or overwhelmed by fire/smoke)."""

    def score(self, done: bool, agent_alive: bool, agent_evacuated: bool, **_) -> float:
        return -10.0 if (done and not agent_alive and not agent_evacuated) else 0.0


class TimeoutPenalty:
    """Penalty for running out of steps while still alive without evacuating.

    Maintains the ordering: success > timeout > death.

    If exits were reachable (not all blocked by fire), the penalty scales with
    remaining health — a healthy agent that timed out demonstrably could have
    evacuated, so it receives a stronger signal (-5 to -8).

    If all exits were fire-blocked, the agent had no path out regardless of
    health, so a flat -5 is applied (not the agent's fault).
    """

    def score(
        self,
        done: bool,
        agent_alive: bool,
        agent_evacuated: bool,
        agent_health: float,
        reachable_exit_count: int,
        **_,
    ) -> float:
        if not (done and agent_alive and not agent_evacuated):
            return 0.0
        if reachable_exit_count == 0:
            return -5.0  # exits were fire-blocked — flat penalty, not the agent's fault
        # Agent had open exits but still timed out — scale by how healthy it was.
        # hp=100 → -8.0,  hp=50 → -6.5,  hp=10 → -5.3
        return -5.0 - 3.0 * (agent_health / 100.0)


class NearMissBonus:
    """Graduated partial credit on DEATH based on closest exit approach.

    Fixes hard-mode reward collapse: when all episodes end in death the flat
    SelfDeathPenalty creates zero gradient. This rubric differentiates "almost
    made it" from "never got close" using the minimum BFS distance reached.

    Formula: max(0.0, 3.0 – 0.5 × min_exit_dist_reached)
      dist=0 → +3.0  (at exit but died — edge case)
      dist=1 → +2.5
      dist=3 → +1.5
      dist=6 → 0.0   (no credit beyond 6 cells away)

    Only fires on DEATH — not on timeout. A timed-out agent was alive and
    had the opportunity to act; softening the TimeoutPenalty here would
    undermine the health-scaled signal from TimeoutPenalty.
    """

    def score(
        self,
        done: bool,
        agent_alive: bool,
        agent_evacuated: bool,
        min_exit_dist_reached: int,
        **_,
    ) -> float:
        # Only award on death (not alive, not evacuated)
        if not done or agent_evacuated or agent_alive:
            return 0.0
        return max(0.0, 3.0 - 0.5 * min_exit_dist_reached)


class TimeBonus:
    """Bonus for escaping quickly — rewards remaining steps when agent evacuates."""

    def score(self, done: bool, agent_evacuated: bool, remaining_steps: int, **_) -> float:
        return 0.05 * remaining_steps if (done and agent_evacuated) else 0.0


# ---------------------------------------------------------------------------
# Convenience factories
# ---------------------------------------------------------------------------

def make_per_step_rubrics():
    return [
        TimeStepPenalty(),
        ProgressReward(),
        ProgressRegressionPenalty(),
        SafeProgressBonus(),
        DangerPenalty(),
        HealthDrainPenalty(),
        StrategicDoorBonus(),
        ExplorationBonus(),
    ]


def make_episode_end_rubrics():
    return [
        SelfSurviveBonus(),
        HealthSurvivalBonus(),
        SelfDeathPenalty(),
        TimeoutPenalty(),
        NearMissBonus(),
        TimeBonus(),
    ]