File size: 2,969 Bytes
7aca0bc | 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 | # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
"""Chess-specific rubrics for reward computation.
Provides ChessWinLossRubric that extends ExponentialDiscountingTrajectoryRubric
to compute temporally-discounted rewards based on chess game outcomes.
See RFC 004 for rubric design: rfcs/004-rubrics.md
"""
from typing import Any, List, Tuple
try:
from openenv.core.rubrics.trajectory import ExponentialDiscountingTrajectoryRubric
except ModuleNotFoundError:
class ExponentialDiscountingTrajectoryRubric:
"""Compatibility fallback when the installed core lacks rubrics."""
def __init__(self, gamma: float = 0.99, intermediate_reward: float = 0.0):
self.gamma = gamma
self.intermediate_reward = intermediate_reward
self._trajectory: List[Tuple[Any, Any]] = []
def __call__(self, action: Any, observation: Any) -> float:
self._trajectory.append((action, observation))
if getattr(observation, "done", False):
return self.score_trajectory(self._trajectory)
return self.intermediate_reward
def reset(self) -> None:
self._trajectory = []
def compute_step_rewards(self) -> List[float]:
if not self._trajectory:
return []
final_score = self.score_trajectory(self._trajectory)
total_steps = len(self._trajectory)
return [
self.gamma ** (total_steps - 1 - step_index) * final_score
for step_index in range(total_steps)
]
class ChessWinLossRubric(ExponentialDiscountingTrajectoryRubric):
"""Score chess game based on win/loss/draw outcome with temporal discounting.
Per-step reward: r_t = gamma^(T-1-t) * R_final
Terminal rewards:
- Win: +1.0
- Loss: -1.0
- Draw: 0.0
Note: Unlike the base class convention of 0.0-1.0 scores, chess uses
-1.0 for losses to match standard chess RL conventions.
Usage:
rubric = ChessWinLossRubric(gamma=0.99)
rubric.reset()
for action, obs in episode:
reward = rubric(action, obs) # 0.0 until done
step_rewards = rubric.compute_step_rewards()
"""
def score_trajectory(self, trajectory: List[Tuple[Any, Any]]) -> float:
"""Score based on game outcome from final observation.
Reads the reward from the terminal observation, which encodes:
+1.0 for agent win, -1.0 for agent loss, 0.0 for draw.
Args:
trajectory: List of (action, observation) tuples.
Returns:
Terminal reward: +1.0 (win), -1.0 (loss), 0.0 (draw).
"""
if not trajectory:
return 0.0
_, final_obs = trajectory[-1]
return getattr(final_obs, "reward", 0.0)
|