diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index e8b8fa283..a543cba2c 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 a831d12bb..034b661b1 100644 --- a/envs/textarena_env/server/environment.py +++ b/envs/textarena_env/server/environment.py @@ -9,11 +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 @@ -42,6 +45,9 @@ _TEXTARENA_IMPORT_ERROR: Exception | None = None _NLTK_DOWNLOADED: bool = False +# 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( r"You submitted \[(\w+)\]\.\nFeedback:\n[^\n]*\n([GYX ]+)" @@ -167,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, @@ -181,7 +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. - self._ta_env.reset(num_players=self.num_players) + self._rng_state: Optional[tuple] = None + self._reset_ta_env() # ------------------------------------------------------------------ # Environment interface @@ -204,7 +212,7 @@ def reset( if hasattr(env, "full_observations"): env.full_observations = {} - self._ta_env.reset(num_players=self.num_players) + self._reset_ta_env(seed) for provider in self._reward_providers: provider.reset() @@ -227,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) @@ -297,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 aa31c451b..8ea50ed03 100644 --- a/tests/envs/test_textarena_environment.py +++ b/tests/envs/test_textarena_environment.py @@ -1,4 +1,7 @@ +import random +import threading from types import SimpleNamespace +from typing import Dict, List import pytest from textarena_env.models import TextArenaAction, TextArenaMessage @@ -83,6 +86,51 @@ def test_textarena_web_playground(): 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" + + def test_wordle_reports_how_the_game_ended(): pytest.importorskip("textarena", reason="textarena not installed") env = TextArenaEnvironment(env_id="Wordle-v0", num_players=1) @@ -176,3 +224,315 @@ def test_sudoku_terminal_invalid_move_reaches_cumulative_client(): assert "invalid" in feedback.lower() assert "out of bounds" in feedback assert obs.messages[0].sender_id == -1 + + +def test_seeded_reset_restores_global_rng_state(): + """A seeded reset must leave the process-global RNG where it found it.""" + pytest.importorskip("textarena", reason="textarena not installed") + env = TextArenaEnvironment(env_id="Wordle-v0", num_players=1) + + random.seed(42) + expected = random.random() + + random.seed(42) + env.reset(seed=1234) + + assert random.random() == expected, "seeded reset leaked into the global RNG" + + +def test_seeded_reset_does_not_perturb_unseeded_session(): + """An unseeded session must get the same episode whether or not another + session performed a seeded reset in between.""" + pytest.importorskip("textarena", reason="textarena not installed") + unseeded = TextArenaEnvironment(env_id="Wordle-v0", num_players=1) + seeded = TextArenaEnvironment(env_id="Wordle-v0", num_players=1) + + random.seed(42) + unseeded.reset() + baseline = _secret_word(unseeded) + + random.seed(42) + seeded.reset(seed=1234) + unseeded.reset() + + assert _secret_word(unseeded) == baseline + + +def test_concurrent_unseeded_reset_cannot_steal_a_seeded_draw(monkeypatch): + """Force another session to reset in the gap between TextArena seeding the + global RNG and drawing the episode. The seeded episode must be unaffected.""" + pytest.importorskip("textarena", reason="textarena not installed") + seeded = TextArenaEnvironment(env_id="Wordle-v0", num_players=1) + other = TextArenaEnvironment(env_id="Wordle-v0", num_players=1) + + seeded.reset(seed=1234) + expected = _secret_word(seeded) + + window_open = threading.Event() + other_done = threading.Event() + original_seed = random.seed + + def seed_then_yield(*args, **kwargs): + original_seed(*args, **kwargs) + window_open.set() + other_done.wait(timeout=0.5) + + monkeypatch.setattr(random, "seed", seed_then_yield) + + def unseeded_reset(): + window_open.wait(timeout=5) + other.reset() + other_done.set() + + thread = threading.Thread(target=unseeded_reset) + thread.start() + seeded.reset(seed=1234) + thread.join(timeout=5) + + assert not thread.is_alive(), "the unseeded reset never completed" + assert _secret_word(seeded) == expected, "another session stole the seeded draw" + + +@pytest.fixture +def fake_textarena(monkeypatch): + """Exercise the wrapper in required CI without optional TextArena/NLTK data.""" + from types import SimpleNamespace + + from textarena_env.server import environment + + class Game: + def reset(self, num_players, seed=None): + random.seed(seed) + self.draw = random.random() + self.state = SimpleNamespace( + turn=0, game_state={}, done=False, game_info={} + ) + + def step(self, message): + self.draw = random.random() + return False, {} + + def get_observation(self): + return 0, [] + + monkeypatch.setattr( + environment, "_TEXTARENA_MODULE", SimpleNamespace(make=lambda **kw: Game()) + ) + rng_state = random.getstate() + yield lambda: TextArenaEnvironment(env_id="test", download_nltk=False) + random.setstate(rng_state) + + +def test_seed_forwarding_without_optional_dependencies(fake_textarena): + env = fake_textarena() + before = random.getstate() + env.reset(seed=0) + assert env._ta_env.draw == random.Random(0).random() + assert random.getstate() == before + + +def test_failed_seeded_reset_restores_rng(fake_textarena, monkeypatch): + env = fake_textarena() + before = random.getstate() + + def fail(**kwargs): + random.seed(kwargs["seed"]) + raise RuntimeError("reset failed") + + monkeypatch.setattr(env._ta_env, "reset", fail) + with pytest.raises(RuntimeError, match="reset failed"): + env.reset(seed=1234) + assert random.getstate() == before + + +@pytest.mark.parametrize("operation", ["reset", "step", "construct"]) +def test_seed_window_excludes_other_session_operations( + fake_textarena, monkeypatch, operation +): + from concurrent.futures import ThreadPoolExecutor + + from textarena_env.server import environment + + seeded = fake_textarena() + other = fake_textarena() + window_open = threading.Event() + attempted = threading.Event() + entered = threading.Event() + release = threading.Event() + lock = threading.Lock() + + class ObservedLock: + def __enter__(self): + if window_open.is_set(): + attempted.set() + lock.acquire() + + def __exit__(self, *args): + lock.release() + + original_reset = seeded._ta_env.reset + + def paused_reset(**kwargs): + original_reset(**kwargs) + window_open.set() + assert release.wait(5), "seed window was never released" + + def draw(*args, **kwargs): + entered.set() + other._ta_env.draw = random.random() + return False, {} + + monkeypatch.setattr(environment, "_RNG_LOCK", ObservedLock()) + monkeypatch.setattr(seeded._ta_env, "reset", paused_reset) + monkeypatch.setattr( + other._ta_env, operation if operation != "construct" else "reset", draw + ) + if operation == "construct": + monkeypatch.setattr( + environment._TEXTARENA_MODULE, "make", lambda **kw: other._ta_env + ) + before = random.getstate() + expected = random.Random() + expected.setstate(before) + expected_draw = expected.random() + + def run_other(): + if operation == "construct": + fake_textarena() + elif operation == "step": + other.step(TextArenaAction(message="draw")) + else: + other.reset() + + with ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(seeded.reset, seed=1234) + try: + assert window_open.wait(5) + second = pool.submit(run_other) + assert attempted.wait(1), "operation did not acquire the RNG lock" + assert not entered.is_set(), "operation entered the seeded RNG window" + finally: + release.set() + first.result(timeout=5) + second.result(timeout=5) + assert other._ta_env.draw == expected_draw + assert random.getstate() == expected.getstate() + + +def _seeded_stream(seed: int, count: int) -> List[float]: + """The draws TextArena would make from ``random.seed(seed)`` onwards.""" + rng = random.Random(seed) + return [rng.random() for _ in range(count)] + + +def test_seeded_trajectory_is_reproducible_across_sessions(fake_textarena): + """Equal seeds and equal actions must give equal draws through step, not + only through reset, and must follow the seed's own stream.""" + first = fake_textarena() + second = fake_textarena() + + draws = [] + for env in (first, second): + env.reset(seed=1234) + trajectory = [env._ta_env.draw] + for _ in range(3): + env.step(TextArenaAction(message="draw")) + trajectory.append(env._ta_env.draw) + draws.append(trajectory) + + assert draws[0] == draws[1], "same seed and actions diverged across sessions" + assert draws[0] == _seeded_stream(1234, 4), "steps left the seeded stream" + + +def test_seeded_trajectory_survives_interleaved_sessions(fake_textarena): + """Other sessions and unrelated draws between a seeded session's calls + must not change that session's trajectory or see its RNG state.""" + seeded = fake_textarena() + unseeded = fake_textarena() + other_seed = fake_textarena() + + seeded.reset(seed=1234) + trajectory = [seeded._ta_env.draw] + for _ in range(3): + unseeded.reset() + unseeded.step(TextArenaAction(message="draw")) + other_seed.reset(seed=7) + other_seed.step(TextArenaAction(message="draw")) + random.random() + + before = random.getstate() + seeded.step(TextArenaAction(message="draw")) + assert random.getstate() == before, "seeded step leaked into the global RNG" + trajectory.append(seeded._ta_env.draw) + + assert trajectory == _seeded_stream(1234, 4) + + +def test_failed_seeded_step_restores_rng_and_continues_stream( + fake_textarena, monkeypatch +): + env = fake_textarena() + env.reset(seed=1234) + original_step = env._ta_env.step + + def fail(message): + original_step(message) + raise RuntimeError("step failed") + + monkeypatch.setattr(env._ta_env, "step", fail) + before = random.getstate() + with pytest.raises(RuntimeError, match="step failed"): + env.step(TextArenaAction(message="draw")) + assert random.getstate() == before + + monkeypatch.setattr(env._ta_env, "step", original_step) + env.step(TextArenaAction(message="draw")) + # reset, failed step, then this step: the failed draw was still consumed. + assert env._ta_env.draw == _seeded_stream(1234, 3)[2] + + +def test_unseeded_reset_returns_session_to_process_rng(fake_textarena): + """An unseeded reset after a seeded episode must draw from, and advance, + the process-global RNG rather than continue the old seeded stream.""" + env = fake_textarena() + env.reset(seed=1234) + env.reset() + + expected = random.Random() + expected.setstate(random.getstate()) + expected_draw = expected.random() + + env.step(TextArenaAction(message="draw")) + + assert env._ta_env.draw == expected_draw + assert random.getstate() == expected.getstate() + + +def _bandit_history(env: TextArenaEnvironment) -> Dict[str, List[float]]: + return env._ta_env.state.game_state["history"] + + +def test_bandit_seeded_trajectory_is_reproducible_across_sessions(): + """Bandit samples its rewards during step. Two sessions with the same seed + and actions must see identical rewards even with a third session in between.""" + pytest.importorskip("textarena", reason="textarena not installed") + first, second, bystander = ( + TextArenaEnvironment(env_id="Bandit-v0", download_nltk=False) for _ in range(3) + ) + actions = ["[red]", "[blue]", "[red]", "[green]", "[red]", "[yellow]"] + + first.reset(seed=1234) + second.reset(seed=1234) + assert ( + first._ta_env.state.game_state["ground_truth"] + == second._ta_env.state.game_state["ground_truth"] + ) + + for message in actions: + first.step(TextArenaAction(message=message)) + bystander.reset() + bystander.step(TextArenaAction(message=message)) + second.step(TextArenaAction(message=message)) + + assert sum(len(v) for v in _bandit_history(first).values()) == len(actions) + assert _bandit_history(first) == _bandit_history(second)