add wordle specific rewards to the environment

Author: burtenshaw

src/envs/textarena_env/__init__.py

@@ -13,12 +13,14 @@
     TextArenaObservation,
     TextArenaState,
 )
+from .rewards import RewardProvider, build_reward_providers
 
 __all__ = [
     "TextArenaEnv",
     "TextArenaAction",
     "TextArenaObservation",
     "TextArenaState",
     "TextArenaMessage",
+    "RewardProvider",
+    "build_reward_providers",
 ]
-

src/envs/textarena_env/rewards.py

@@ -0,0 +1,133 @@
+"""Reward provider utilities for TextArena environments."""
+
+from __future__ import annotations
+
+import re
+from typing import Dict, List, Protocol, Tuple
+
+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)
+
+    segments = [
+        segment.strip() for segment in feedback.split("\n\n") if segment.strip()
+    ]
+    if not segments:
+        return (0, 0)
+
+    latest_segment = segments[-1]
+    lines = [line.strip() for line in latest_segment.splitlines() if line.strip()]
+    latest_line = lines[-1] if lines else latest_segment
+
+    green_count = latest_line.count("G")
+    yellow_count = latest_line.count("Y")
+    return (green_count, yellow_count)
+
+
+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",
+]

src/envs/textarena_env/server/environment.py

@@ -17,6 +17,7 @@
 from core.env_server.interfaces import Environment
 
 from ..models import TextArenaAction, TextArenaMessage, TextArenaObservation, TextArenaState
+from ..rewards import RewardProvider, build_reward_providers
 
 
 _TEXTARENA_MODULE: Any | None = None
@@ -84,18 +85,25 @@ def __init__(
             max_turns=max_turns,
         )
 
+        self._reward_providers: List[RewardProvider] = build_reward_providers(env_id)
+        self._last_reward_signals: Dict[str, float] = {}
+
     # ------------------------------------------------------------------
     # Environment interface
     # ------------------------------------------------------------------
     def reset(self) -> TextArenaObservation:
         self._ta_env.reset(num_players=self.num_players)
 
+        for provider in self._reward_providers:
+            provider.reset()
+
         self._state.episode_id = str(uuid4())
         self._state.step_count = 0
         self._state.turn = 0
         self._state.last_reward = 0.0
         self._state.last_info = {}
         self._state.raw_state = self._snapshot_state()
+        self._last_reward_signals = {}
 
         observation = self._build_observation()
         observation.reward = 0.0
@@ -119,6 +127,14 @@ def step(self, action: TextArenaAction) -> TextArenaObservation:  # type: ignore
         reward = self._extract_reward()
         observation.reward = reward
         self._state.last_reward = reward
+
+        reward_signals = self._compute_reward_signals(action=action, observation=observation)
+        if reward_signals:
+            observation.info.setdefault("reward_signals", {}).update(reward_signals)
+            observation.metadata.setdefault("reward_signals", {}).update(reward_signals)
+        self._last_reward_signals = reward_signals
+        if reward_signals:
+            self._state.last_info = {**(self._state.last_info or {}), "reward_signals": reward_signals}
         self._state.raw_state = self._snapshot_state()
 
         return observation
@@ -214,5 +230,23 @@ def _snapshot_state(self) -> Dict[str, Any]:
             "game_info": getattr(state, "game_info", {}),
             "step_info": getattr(state, "step_info", {}),
         }
+        if self._last_reward_signals:
+            snapshot["reward_signals"] = dict(self._last_reward_signals)
         return snapshot
 
+    def _compute_reward_signals(
+        self, *, action: TextArenaAction, observation: TextArenaObservation
+    ) -> Dict[str, float]:
+        if not self._reward_providers:
+            return {}
+
+        aggregated: Dict[str, float] = {}
+        for provider in self._reward_providers:
+            try:
+                result = provider.compute(action=action, observation=observation)
+            except Exception:  # pragma: no cover - defensive
+                continue
+            for key, value in result.items():
+                aggregated[key] = float(value)
+        return aggregated
+

src/envs/textarena_env/server/run_local.sh

@@ -1,6 +1,7 @@
-TEXTARENA_ENV_ID="Wordle-v0" TEXTARENA_NUM_PLAYERS=2
+export TEXTARENA_ENV_ID="Wordle-v0" 
+export TEXTARENA_NUM_PLAYERS=1
 
 # Run the server
-exec uvicorn envs.textarena_env.server.app:app --host 0.0.0.0 --port 8000
+exec uvicorn envs.textarena_env.server.app:app --host 0.0.0.0 --port 8001