From 90b7a6b22598473a95f50547f5f906240b9af87f Mon Sep 17 00:00:00 2001 From: akashjainn Date: Wed, 14 Oct 2026 13:44:23 -0400 Subject: [PATCH 1/2] fix(textarena_env): forward reset(seed=...) so episodes are reproducible TextArenaEnvironment.reset accepted a seed and never passed it to the wrapped TextArena env, so seeded runs were not reproducible. TextArena seeds and draws from the process-global random module, and the server runs several sessions per process, so each reset's seed-and-draw runs under a process lock to keep it atomic. A seeded reset still leaves the global RNG in a known state for other sessions; the proper fix for that is a per-instance RNG in TextArena. Fixes #1102 Co-Authored-By: Claude Fable 5.1 --- envs/textarena_env/server/environment.py | 12 +++++-- tests/envs/test_textarena_environment.py | 45 ++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/envs/textarena_env/server/environment.py b/envs/textarena_env/server/environment.py index f8f43307a..66242f268 100644 --- a/envs/textarena_env/server/environment.py +++ b/envs/textarena_env/server/environment.py @@ -12,6 +12,7 @@ import re import subprocess import sys +import threading import urllib.request from typing import Any, Dict, Iterable, List, Optional from uuid import uuid4 @@ -42,6 +43,11 @@ _TEXTARENA_IMPORT_ERROR: Exception | None = None _NLTK_DOWNLOADED: bool = False +# TextArena seeds and draws from the process-global ``random`` module, so a +# seeded reset in one session can race with a draw in another. The lock keeps +# each reset's seed-and-draw atomic. +_RESET_LOCK = threading.Lock() + # TextArena's Wordle feedback, e.g. "You submitted [crane].\nFeedback:\nC R A N E\nG X Y X G". _WORDLE_FEEDBACK = re.compile( r"You submitted \[(\w+)\]\.\nFeedback:\n[^\n]*\n([GYX ]+)" @@ -178,7 +184,8 @@ def __init__( # Initialize environment state - TextArena envs require reset() to be called # before step() can be used, as the internal state object isn't created until reset. # This ensures the environment is always in a valid state after construction. - self._ta_env.reset(num_players=self.num_players) + with _RESET_LOCK: + self._ta_env.reset(num_players=self.num_players) # ------------------------------------------------------------------ # Environment interface @@ -201,7 +208,8 @@ def reset( if hasattr(env, "full_observations"): env.full_observations = {} - self._ta_env.reset(num_players=self.num_players) + with _RESET_LOCK: + self._ta_env.reset(num_players=self.num_players, seed=seed) for provider in self._reward_providers: provider.reset() diff --git a/tests/envs/test_textarena_environment.py b/tests/envs/test_textarena_environment.py index 5552ccbed..45b89fd8c 100644 --- a/tests/envs/test_textarena_environment.py +++ b/tests/envs/test_textarena_environment.py @@ -79,3 +79,48 @@ def test_textarena_web_playground(): assert "#" not in board # theme colours only, so it reads in dark mode assert env.render_web({"messages": []}) is None assert env.render_web({}) is None + + +def _secret_word(env: TextArenaEnvironment) -> str: + # Wordle's prompt is static instructions, so the puzzle has to be read from + # TextArena's own game state to tell episodes apart. + return env._ta_env.state.game_state["secret_word"] + + +def test_reset_seed_is_forwarded_to_textarena(): + pytest.importorskip("textarena", reason="textarena not installed") + env = TextArenaEnvironment(env_id="Wordle-v0", num_players=1) + + seeded = [] + for _ in range(3): + env.reset(seed=1234) + seeded.append(_secret_word(env)) + + env.reset(seed=999) + other_seed = _secret_word(env) + + assert len(set(seeded)) == 1, f"seed=1234 produced {sorted(set(seeded))}" + assert other_seed != seeded[0], "a different seed produced the same episode" + + +def test_reset_seed_is_reproducible_across_instances(): + pytest.importorskip("textarena", reason="textarena not installed") + first = TextArenaEnvironment(env_id="Wordle-v0", num_players=1) + first.reset(seed=1234) + + second = TextArenaEnvironment(env_id="Wordle-v0", num_players=1) + second.reset(seed=1234) + + assert _secret_word(first) == _secret_word(second) + + +def test_reset_without_seed_still_varies(): + pytest.importorskip("textarena", reason="textarena not installed") + env = TextArenaEnvironment(env_id="Wordle-v0", num_players=1) + + words = set() + for _ in range(8): + env.reset() + words.add(_secret_word(env)) + + assert len(words) > 1, "unseeded resets should not be deterministic" From b08a223cefefa153dad6fa1bce7e3ad8b003988c Mon Sep 17 00:00:00 2001 From: burtenshaw Date: Mon, 26 Oct 2026 15:39:20 +0100 Subject: [PATCH 2/2] fix(textarena): preserve seeded trajectories across sessions (#1102) --- .github/workflows/test.yml | 2 +- envs/textarena_env/server/environment.py | 67 ++++++++++++++++++++---- tests/envs/test_textarena_environment.py | 7 +-- 3 files changed, 61 insertions(+), 15 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 0f9ac4fb7..e3fe3450e 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -89,7 +89,7 @@ jobs: run: | uv sync --all-extras uv pip install pytest pytest-asyncio numpy nltk smolagents python-chess moonfish \ - langdetect "math-verify[antlr4_13_2]" datasets open_spiel + langdetect "math-verify[antlr4_13_2]" datasets open_spiel "textarena==0.7.4" - name: Run tests run: | diff --git a/envs/textarena_env/server/environment.py b/envs/textarena_env/server/environment.py index 9be285eff..034b661b1 100644 --- a/envs/textarena_env/server/environment.py +++ b/envs/textarena_env/server/environment.py @@ -9,12 +9,14 @@ from __future__ import annotations import os +import random import re import subprocess import sys import threading import urllib.request -from typing import Any, Dict, Iterable, List, Optional +from contextlib import contextmanager +from typing import Any, Dict, Iterable, Iterator, List, Optional from uuid import uuid4 from openenv.core.env_server.interfaces import Environment @@ -43,10 +45,8 @@ _TEXTARENA_IMPORT_ERROR: Exception | None = None _NLTK_DOWNLOADED: bool = False -# TextArena seeds and draws from the process-global ``random`` module, so a -# seeded reset in one session can race with a draw in another. The lock keeps -# each reset's seed-and-draw atomic. -_RESET_LOCK = threading.Lock() +# Serialize TextArena RNG access; seeded sessions retain their own stream. +_RNG_LOCK = threading.Lock() # TextArena's Wordle feedback, e.g. "You submitted [crane].\nFeedback:\nC R A N E\nG X Y X G". _WORDLE_FEEDBACK = re.compile( @@ -173,7 +173,8 @@ def __init__( self.max_turns = max_turns self._env_kwargs = env_kwargs or {} - self._ta_env = ta.make(env_id=env_id, **self._env_kwargs) + with _RNG_LOCK: + self._ta_env = ta.make(env_id=env_id, **self._env_kwargs) self._state = TextArenaState( env_id=env_id, @@ -187,8 +188,8 @@ def __init__( # Initialize environment state - TextArena envs require reset() to be called # before step() can be used, as the internal state object isn't created until reset. # This ensures the environment is always in a valid state after construction. - with _RESET_LOCK: - self._ta_env.reset(num_players=self.num_players) + self._rng_state: Optional[tuple] = None + self._reset_ta_env() # ------------------------------------------------------------------ # Environment interface @@ -211,8 +212,7 @@ def reset( if hasattr(env, "full_observations"): env.full_observations = {} - with _RESET_LOCK: - self._ta_env.reset(num_players=self.num_players, seed=seed) + self._reset_ta_env(seed) for provider in self._reward_providers: provider.reset() @@ -235,7 +235,8 @@ def step(self, action: TextArenaAction) -> TextArenaObservation: # type: ignore if not isinstance(action, TextArenaAction): raise TypeError(f"Expected TextArenaAction, received {type(action)!r}") - done, info = self._ta_env.step(action.message) + with self._session_rng(): + done, info = self._ta_env.step(action.message) self._state.step_count += 1 self._state.turn = getattr(self._ta_env.state, "turn", self._state.turn + 1) @@ -305,6 +306,50 @@ def state(self) -> TextArenaState: # ------------------------------------------------------------------ # Helpers # ------------------------------------------------------------------ + def _reset_ta_env(self, seed: Optional[int] = None) -> None: + """Reset the wrapped game without disturbing other sessions. + + TextArena implements ``reset(seed=...)`` by calling ``random.seed`` on the + process-global RNG and then drawing the episode from it. The lock keeps + seed and draw atomic across concurrent sessions. A seeded reset keeps the + RNG state it produced on the session, so later ``step`` calls continue the + same stream, and restores the prior process state so a seeded reset does + not make unseeded sessions predictable. + """ + with _RNG_LOCK: + self._rng_state = None + if seed is None: + self._ta_env.reset(num_players=self.num_players) + return + process_state = random.getstate() + try: + self._ta_env.reset(num_players=self.num_players, seed=seed) + self._rng_state = random.getstate() + finally: + random.setstate(process_state) + + @contextmanager + def _session_rng(self) -> Iterator[None]: + """Run a wrapped call under this session's RNG stream. + + Unseeded sessions draw from the process-global RNG as before. A seeded + session swaps its saved state in, runs the call, saves the advanced + state back, and restores the process state on exit, including when the + call raises, so equal seeds and equal actions give equal trajectories + regardless of what other sessions do in between. + """ + with _RNG_LOCK: + if self._rng_state is None: + yield + return + process_state = random.getstate() + random.setstate(self._rng_state) + try: + yield + finally: + self._rng_state = random.getstate() + random.setstate(process_state) + def _build_observation(self) -> TextArenaObservation: player_id, messages = self._ta_env.get_observation() diff --git a/tests/envs/test_textarena_environment.py b/tests/envs/test_textarena_environment.py index feeb8397c..8ea50ed03 100644 --- a/tests/envs/test_textarena_environment.py +++ b/tests/envs/test_textarena_environment.py @@ -1,8 +1,7 @@ import random import threading -from typing import Dict, List - from types import SimpleNamespace +from typing import Dict, List import pytest from textarena_env.models import TextArenaAction, TextArenaMessage @@ -305,7 +304,9 @@ class Game: def reset(self, num_players, seed=None): random.seed(seed) self.draw = random.random() - self.state = SimpleNamespace(turn=0, game_state={}) + self.state = SimpleNamespace( + turn=0, game_state={}, done=False, game_info={} + ) def step(self, message): self.draw = random.random()