Open a new client per game in the Wordle GRPO script

docs/source/tutorials/sudoku-grpo.md

@@ -166,7 +166,7 @@ def reset(self, **kwargs):
 
 The old session is released asynchronously, which is why `MAX_CONCURRENT_ENVS` is 96 rather than 64.
 
-Since [#1536](https://github.com/openenv/OpenEnv/pull/1536), the client no longer reconnects silently in the middle of an episode. It raises `SessionLostError` and the next `reset()` starts a new session. A dropped game then fails instead of being played on the wrong board, and a new connection per game still keeps the drops away from training.
+Since OpenEnv 0.12.0 ([#1536](https://github.com/openenv/OpenEnv/pull/1536)), the client no longer reconnects silently in the middle of an episode. It raises `SessionLostError` and the next `reset()` starts a new session. A dropped game then fails instead of being played on the wrong board, and a new connection per game still keeps the drops away from training.
 
 ## 8. Evaluate and make the GIF
 

docs/source/tutorials/wordle-grpo.md

@@ -117,7 +117,7 @@ hf jobs uv run --flavor a100-large --timeout 4h --secrets HF_TOKEN grpo.py \
     --env-url https://<user>-wordle.hf.space
 ```
 
-Every game is one environment session, so the environment must accept 64 at once. [Duplicate the `openenv/wordle` Space](https://huggingface.co/spaces/openenv/wordle?duplicate=true) and set the `MAX_CONCURRENT_ENVS=64` variable on your copy, rather than training against the shared Space. A private copy works too, since the script sends your token.
+Every game is one environment session and GRPO plays 64 at once. Each game opens a new session before the server has released the previous one, so the environment needs some room above 64. [Duplicate the `openenv/wordle` Space](https://huggingface.co/spaces/openenv/wordle?duplicate=true) and set the `MAX_CONCURRENT_ENVS=96` variable on your copy, rather than training against the shared Space. A private copy works too, since the script sends your token.
 
 ![Reward during GRPO](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/openenv/wordle-grpo-reward.png)
 

examples/sudoku/grpo.py

@@ -56,7 +56,6 @@ def main():
     parser.add_argument("--trackio-space-id", default=None)
     args = parser.parse_args()
     env_url = args.env_url
-    # A private Space needs your token. A public one ignores it.
     token = get_token()
     headers = {"Authorization": f"Bearer {token}"} if token else None
 
@@ -65,8 +64,7 @@ def __init__(self):
             self.client = TextArenaEnv(base_url=env_url, headers=headers).sync()
 
         def reset(self, **kwargs) -> str | None:
-            # A new connection per game. The Space closes WebSockets after about 2 hours, and a game cut in the middle
-            # is lost.
+            # One connection per game, so the Space's periodic WebSocket drops don't land mid-game.
             self.client.close()
             self.client = TextArenaEnv(base_url=env_url, headers=headers).sync()
             result = self.client.reset()
@@ -105,7 +103,8 @@ def place(self, row: int, col: int, number: int) -> str:
             self.solved = result.done and result.reward >= 1.0
             return feedback.split("Board state: \n")[-1]
 
-    # 4 prompts x 16 games per step. Each game is one environment session, so the Space must allow 64 at once.
+    # 4 prompts x 16 games per step, one environment session each. MAX_CONCURRENT_ENVS=96 leaves room for the sessions
+    # still being closed between games.
     # One game per forward pass: with games up to 16k tokens, 8 at once ran out of memory on an A100 80GB.
     prompts_per_step, num_generations, per_device_batch_size = 4, 16, 1
     dataset = Dataset.from_dict(
@@ -124,7 +123,7 @@ def place(self, row: int, col: int, number: int) -> str:
             output_dir=args.hub_model_id.split("/")[-1],
             max_steps=args.max_steps,
             learning_rate=1e-6,
-            # Constant: in the Wordle run, the reward was still rising when a linear decay ran out.
+            # Constant, since with a linear decay to 0 the Wordle reward was still rising when the learning rate ran out.
             lr_scheduler_type="constant",
             num_generations=num_generations,
             per_device_train_batch_size=per_device_batch_size,
@@ -134,7 +133,6 @@ def place(self, row: int, col: int, number: int) -> str:
             gradient_checkpointing=True,
             bf16=True,
             chat_template_kwargs={"enable_thinking": True},
-            # Every move returns the whole grid: a 26-cell game is about 13k tokens.
             max_completion_length=16384,
             mask_truncated_completions=True,
             # Multi-turn rollouts: clip per-token ratios instead of masking whole sequences.

examples/wordle/README.md

@@ -22,4 +22,4 @@ uv run evaluate.py --model <user>/Qwen3-1.7B-Wordle-SFT-GRPO --thinking --max-ne
     --secrets "$(cat secrets30.txt)" --out grpo.jsonl
 ```
 
-GRPO plays 64 games at once, one environment session each. Duplicate the `openenv/wordle` Space and set `MAX_CONCURRENT_ENVS=64` on your copy before training.
+GRPO plays 64 games at once, one environment session each, plus the sessions still being closed between games. Duplicate the `openenv/wordle` Space and set `MAX_CONCURRENT_ENVS=96` on your copy before training.

examples/wordle/grpo.py

@@ -70,15 +70,17 @@ def main():
     parser.add_argument("--trackio-space-id", default=None)
     args = parser.parse_args()
     env_url = args.env_url
-    # A private Space needs your token. A public one ignores it.
     token = get_token()
     headers = {"Authorization": f"Bearer {token}"} if token else None
 
     class WordleEnv:
         def __init__(self):
-            self.client = TextArenaEnv(base_url=env_url, headers=headers)
+            self.client = TextArenaEnv(base_url=env_url, headers=headers).sync()
 
         def reset(self, **kwargs) -> str | None:
+            # One connection per game, so the Space's periodic WebSocket drops don't land mid-game.
+            self.client.close()
+            self.client = TextArenaEnv(base_url=env_url, headers=headers).sync()
             result = self.client.reset()
             # The game returns cumulative feedback each turn, so keep the full text and return only the new part.
             self._last_full_feedback = result.observation.messages[0].content
@@ -116,7 +118,8 @@ def guess(self, guess: str) -> str:
             self.done = result.done
             return feedback
 
-    # 4 prompts x 16 games per step. Each game is one environment session, so the Space must allow 64 at once.
+    # 4 prompts x 16 games per step, one environment session each. MAX_CONCURRENT_ENVS=96 leaves room for the sessions
+    # still being closed between games.
     prompts_per_step, num_generations, per_device_batch_size = 4, 16, 8
     dataset = Dataset.from_dict(
         {
@@ -134,7 +137,7 @@ def guess(self, guess: str) -> str:
             output_dir=args.hub_model_id.split("/")[-1],
             max_steps=args.max_steps,
             learning_rate=1e-6,
-            # Constant: with a linear decay to 0 the reward was still rising when the learning rate ran out.
+            # Constant, since with a linear decay to 0 the Wordle reward was still rising when the learning rate ran out.
             lr_scheduler_type="constant",
             num_generations=num_generations,
             per_device_train_batch_size=per_device_batch_size,

tests/scripts/test_wordle_training.py

@@ -86,6 +86,12 @@ def result(self, *, done=False, reward=0.0):
             reward=reward,
         )
 
+    def sync(self):
+        return self
+
+    def close(self):
+        pass
+
     def reset(self):
         return self.result()