| """Reward provider utilities for TextArena environments.""" |
|
|
| from __future__ import annotations |
|
|
| import re |
| from typing import Dict, List, Protocol, Tuple |
|
|
| try: |
| from textarena_env.models import TextArenaAction, TextArenaObservation |
| except ImportError: |
| from models import TextArenaAction, TextArenaObservation |
|
|
|
|
| class RewardProvider(Protocol): |
| """Interface for computing auxiliary reward signals.""" |
|
|
| def reset(self) -> None: |
| """Clear any internal state before a new episode.""" |
|
|
| def compute( |
| self, *, action: TextArenaAction, observation: TextArenaObservation |
| ) -> Dict[str, float]: |
| """Return a mapping of reward names to float values for the step.""" |
|
|
|
|
| def build_reward_providers(env_id: str) -> List[RewardProvider]: |
| """Instantiate reward providers appropriate for the given environment.""" |
|
|
| providers: List[RewardProvider] = [] |
| if env_id == "Wordle-v0": |
| providers.append(_WordleRewardProvider()) |
| return providers |
|
|
|
|
| _WORDLE_GUESS_PATTERN = re.compile(r"\[[A-Za-z]{5}\]") |
|
|
|
|
| def extract_guess(text: str) -> str: |
| """Normalize a Wordle guess string from arbitrary text.""" |
|
|
| match = _WORDLE_GUESS_PATTERN.search(text) |
| if match: |
| return match.group(0).lower() |
|
|
| cleaned = re.sub(r"[^a-z]", "", text.lower()) |
| if len(cleaned) >= 5: |
| return f"[{cleaned[:5]}]" |
| return "[dunno]" |
|
|
|
|
| def extract_wordle_feedback(observation: TextArenaObservation) -> str: |
| """Pull the latest feedback text from a Wordle observation.""" |
|
|
| for message in reversed(observation.messages): |
| content = message.content.strip() |
| if "Feedback:" in content: |
| return content.split("Feedback:", 1)[-1].strip() |
| return "" |
|
|
|
|
| def extract_feedback_counts(feedback: str) -> Tuple[int, int]: |
| """Return counts of green (G) and yellow (Y) markers from feedback.""" |
|
|
| if not feedback: |
| return (0, 0) |
|
|
| lines = [line.strip() for line in feedback.split("\n") if line.strip()] |
| if len(lines) < 2: |
| return (0, 0) |
|
|
| for line in reversed(lines): |
| normalized = line.replace(" ", "") |
| if normalized and all(c in "GYX" for c in normalized): |
| green = normalized.count("G") |
| yellow = normalized.count("Y") |
| return (green, yellow) |
|
|
| return (0, 0) |
|
|
|
|
| class _WordleRewardProvider: |
| """Reward provider that mirrors the GRPO Wordle heuristics.""" |
|
|
| SIGNAL_MAP = { |
| "greens": "wordle.greens", |
| "yellows": "wordle.yellows", |
| "repetitions": "wordle.repetitions", |
| "correct": "wordle.correct", |
| } |
|
|
| def __init__(self) -> None: |
| self._guess_history: Dict[str, int] = {} |
|
|
| def reset(self) -> None: |
| self._guess_history.clear() |
|
|
| def compute( |
| self, *, action: TextArenaAction, observation: TextArenaObservation |
| ) -> Dict[str, float]: |
| guess = extract_guess(action.message) |
| feedback = extract_wordle_feedback(observation) |
|
|
| normalized_guess = guess if guess and guess != "[dunno]" else "" |
| previous_occurrences = ( |
| self._guess_history.get(normalized_guess, 0) if normalized_guess else 0 |
| ) |
|
|
| green_score = 0.0 |
| yellow_score = 0.0 |
| if feedback: |
| green_count, yellow_count = extract_feedback_counts(feedback) |
| green_score = green_count / 5.0 |
| yellow_score = yellow_count / 5.0 |
|
|
| repetition_score = 1.0 - previous_occurrences |
| correct_score = float(observation.reward or 0.0) |
|
|
| if normalized_guess: |
| self._guess_history[normalized_guess] = previous_occurrences + 1 |
|
|
| return { |
| self.SIGNAL_MAP["greens"]: float(green_score), |
| self.SIGNAL_MAP["yellows"]: float(yellow_score), |
| self.SIGNAL_MAP["repetitions"]: float(repetition_score), |
| self.SIGNAL_MAP["correct"]: float(correct_score), |
| } |
|
|
|
|
| __all__ = [ |
| "RewardProvider", |
| "build_reward_providers", |
| "extract_feedback_counts", |
| "extract_guess", |
| "extract_wordle_feedback", |
| ] |
|
|
