Spaces:
Sleeping
Sleeping
File size: 13,393 Bytes
443c22e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 | """
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.1 agent moved closer to nearest unblocked exit (BFS distance)
ProgressRegressionPenalty -0.05 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.
"""
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.1 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.1 for progress
and –0.05 for regression, creating a clean two-sided gradient.
"""
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.05 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.
Lighter than SelfDeathPenalty (–5 vs –10) to maintain the ordering:
success > timeout > death. Creates pressure to act decisively.
"""
def score(self, done: bool, agent_alive: bool, agent_evacuated: bool, **_) -> float:
return -5.0 if (done and agent_alive and not agent_evacuated) else 0.0
class NearMissBonus:
"""Graduated partial credit on failure based on closest exit approach.
Fixes hard-mode reward collapse: when all episodes end in failure the
flat SelfDeathPenalty creates zero gradient. This rubric differentiates
"almost made it" from "never got close" using the minimum BFS distance
the agent reached during the episode.
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 failure (death OR timeout), not on success.
"""
def score(
self,
done: bool,
agent_evacuated: bool,
min_exit_dist_reached: int,
**_,
) -> float:
if not done or agent_evacuated:
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(),
]
|