Add distributed FrozenLake GRPO recipe to experimental examples.

examples/frozenlake/env.py

@@ -176,9 +176,9 @@ def __init__(
 
     desc = kwargs.pop("desc", None)
     is_slippery = kwargs.pop("is_slippery", False)
-    self.seed = entry["seed"].item() if "seed" in entry else 42
-    self.size = entry["size"].item() if "size" in entry else 8
-    self.p = entry["p"].item() if "p" in entry else 0.8
+    self.seed = int(entry.get("seed", 42))
+    self.size = int(entry.get("size", 8))
+    self.p = float(entry.get("p", 0.8))
 
     if desc is None:
       random_map, goal_position = generate_random_map(

tests/experimental/examples/frozenlake_dist/frozenlake_test.py

@@ -0,0 +1,199 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the self-contained distributed FrozenLake recipe."""
+
+from pathlib import Path
+import sys
+import unittest
+
+from absl.testing import absltest
+
+if "pytest" in sys.modules:
+  try:
+    import pytest  # pylint: disable=g-import-not-at-top
+
+    pytest.importorskip("gymnasium")
+  except ImportError:
+    pass
+
+try:
+  import gymnasium  # pylint: disable=unused-import,g-import-not-at-top
+except ImportError:
+  raise unittest.SkipTest("gymnasium is not installed")  # pylint: disable=raise-missing-from
+
+# pylint: disable=g-import-not-at-top
+from examples.frozenlake import agent as frozenlake_agent
+from examples.frozenlake import env as frozenlake_env
+from tunix.experimental.examples.common import models
+from tunix.experimental.examples.frozenlake_dist import frozenlake
+from tunix.experimental.examples.frozenlake_dist import run_frozenlake_dist
+from tunix.experimental.rl.agentic import registry
+# pylint: enable=g-import-not-at-top
+
+
+class FrozenLakeDistTest(absltest.TestCase):
+
+  def test_registered_components(self):
+    self.assertIs(
+        registry.ENV_REGISTRY.get(frozenlake.FROZENLAKE_ENV_NAME),
+        frozenlake.FrozenLakeEnv,
+    )
+    self.assertIs(
+        registry.AGENT_REGISTRY.get(frozenlake.FROZENLAKE_AGENT_NAME),
+        frozenlake.FrozenLakeAgent,
+    )
+    self.assertIs(frozenlake.FrozenLakeEnv, frozenlake_env.FrozenLakeEnv)
+    self.assertIs(frozenlake.FrozenLakeAgent, frozenlake_agent.FrozenLakeAgent)
+
+  def test_package_reuses_recipe_agent_and_env(self):
+    package_dir = Path(frozenlake.__file__).parent
+    source = Path(frozenlake.__file__).read_text(encoding="utf-8")
+    self.assertIn("from examples.frozenlake import agent", source)
+    self.assertIn("from examples.frozenlake import env", source)
+    self.assertNotIn("class FrozenLakeAgent", source)
+    self.assertNotIn("class FrozenLakeEnv", source)
+
+  def test_dataset_is_deterministic_and_uses_reference_ranges(self):
+    first = frozenlake.create_dataset(size=5, seed=123)
+    second = frozenlake.create_dataset(size=5, seed=123)
+
+    self.assertEqual(first, second)
+    self.assertLen(first, 5)
+    for entry in first:
+      self.assertBetween(entry["seed"], 0, 99999)
+      self.assertBetween(entry["size"], 2, 9)
+      self.assertGreaterEqual(entry["p"], 0.6)
+      self.assertLess(entry["p"], 0.85)
+
+  def test_generated_map_is_reproducible_and_reachable(self):
+    previous_max_steps = frozenlake_env.MAX_STEPS
+    frozenlake_env.MAX_STEPS = 8
+    self.addCleanup(setattr, frozenlake_env, "MAX_STEPS", previous_max_steps)
+    first_map, first_goal = frozenlake_env.generate_random_map(
+        size=5, p=0.8, seed=7
+    )
+    second_map, second_goal = frozenlake_env.generate_random_map(
+        size=5, p=0.8, seed=7
+    )
+
+    self.assertEqual(first_map, second_map)
+    self.assertEqual(first_goal, second_goal)
+    self.assertTrue(
+        frozenlake_env.is_valid([list(row) for row in first_map], max_size=5)
+    )
+
+  def test_environment_reaches_goal(self):
+    env = frozenlake.FrozenLakeEnv(
+        entry={"seed": 42, "size": 2, "p": 0.8},
+        desc=["SF", "FG"],
+        is_slippery=False,
+        max_steps=2,
+    )
+    observation, _ = env.reset()
+    self.assertIn("P", observation)
+
+    _, reward, done, _ = env.step("3")
+    self.assertEqual(reward, 0.0)
+    self.assertFalse(done)
+    observation, reward, done, _ = env.step("2")
+    self.assertEqual(reward, 1.0)
+    self.assertTrue(done)
+    self.assertIn("√", observation)
+
+  def test_agent_parses_last_fenced_direction(self):
+    agent = frozenlake.FrozenLakeAgent(use_multistep_prompt=False)
+    agent.update_from_env("P _\nO G", 0.0, False)
+
+    action = agent.update_from_model(
+        "I first considered ```Up```, but the answer is ```Right```."
+    )
+
+    self.assertEqual(action.action, "3")
+    self.assertEqual(agent.trajectory.steps[-1].action, "3")
+
+  def test_prompt_item_carries_distributed_env_and_agent_config(self):
+    item = frozenlake.build_prompt_item(
+        entry={"seed": 1, "size": 4, "p": 0.75},
+        prompt_idx=2,
+        max_turns=8,
+        max_response_length=2048,
+        episode_timeout_secs=600,
+        temperature=0.7,
+        top_p=1.0,
+        top_k=0,
+        is_slippery=False,
+        use_multistep_prompt=True,
+    )
+
+    self.assertEqual(item["prompt_id"], "frozenlake_2")
+    self.assertEqual(item["max_turns"], 8)
+    self.assertEqual(item["generation_kwargs"]["temperature"], 0.7)
+    self.assertEqual(item["metadata"]["env_config"]["max_steps"], 8)
+    self.assertEqual(
+        item["metadata"]["agent_config"], {"use_multistep_prompt": True}
+    )
+
+  def test_prompt_iterator_emits_one_full_batch_per_step(self):
+    items = list(
+        frozenlake.iter_prompt_items(
+            dataset=[{"seed": 1, "size": 2, "p": 0.8}],
+            max_steps=2,
+            batch_size=3,
+            max_turns=4,
+            max_response_length=64,
+            episode_timeout_secs=30,
+            temperature=0.7,
+            top_p=1.0,
+            top_k=0,
+            is_slippery=False,
+            use_multistep_prompt=True,
+        )
+    )
+    self.assertLen(items, 6)
+    self.assertLen({item["prompt_id"] for item in items}, 6)
+
+  def test_recipe_defaults_match_reference(self):
+    args = run_frozenlake_dist._parse_args([])
+    self.assertEqual(args.model_id, "Qwen/Qwen3-8B")
+    self.assertEqual(args.batch_size, 64)
+    self.assertEqual(args.mini_batch_size, 64)
+    self.assertEqual(args.num_generations, 8)
+    self.assertEqual(args.max_turns, 8)
+    self.assertEqual(args.epsilon, 0.003)
+    self.assertEqual(args.epsilon_high, 0.005)
+    self.assertEqual(args.loss_algo, "gspo-token")
+    self.assertEqual(args.advantage_estimator, "rloo")
+
+  def test_qwen3_8b_supported_by_distributed_workers(self):
+    config = models._qwen3_config("Qwen3-8B")
+    self.assertEqual(config.embed_dim, 4096)
+    self.assertEqual(config.num_layers, 36)
+
+  def test_launcher_uses_frozenlake_registry(self):
+    launcher = (Path(frozenlake.__file__).parent / "launcher.sh").read_text(
+        encoding="utf-8"
+    )
+    self.assertIn(
+        "--registry_module=tunix.experimental.examples.frozenlake_dist.frozenlake",
+        launcher,
+    )
+    self.assertIn("--env_name=frozenlake_env", launcher)
+    self.assertIn("--agent_name=frozenlake_agent", launcher)
+    self.assertEqual(launcher.count('--mini_batch_size="$MINI_BATCH_SIZE"'), 2)
+    self.assertEqual(launcher.count('--num_generations="$NUM_GENERATIONS"'), 2)
+
+
+if __name__ == "__main__":
+  absltest.main()

tunix/experimental/examples/common/models.py

@@ -43,6 +43,8 @@ def _qwen3_config(model_name: str) -> qwen3_model_lib.ModelConfig:
     config = qwen3_model_lib.ModelConfig.qwen3_0p6b()
   elif "1.7b" in normalized or "1p7b" in normalized:
     config = qwen3_model_lib.ModelConfig.qwen3_1p7b()
+  elif "8b" in normalized:
+    config = qwen3_model_lib.ModelConfig.qwen3_8b()
   elif "32b" in normalized:
     config = qwen3_model_lib.ModelConfig.qwen3_32b()
   else:

tunix/experimental/examples/frozenlake_dist/README.md

@@ -0,0 +1,56 @@
+# Distributed FrozenLake GRPO Recipe
+
+This directory ports `examples/frozenlake/train_frozenlake_qwen3.py` to the
+experimental distributed RL stack. The control plane runs on CPU, while the
+generic distributed trainer and rollout workers own model execution.
+
+The distributed module keeps only its in-memory dataset and request wiring. It
+directly registers and reuses `examples/frozenlake/agent.py` and
+`examples/frozenlake/env.py`, avoiding a second copy of the recipe behavior.
+
+The defaults preserve the reference Qwen3 recipe: Qwen3-8B, 64 prompt groups
+per full step, 64 prompt groups per optimizer update, 8 generations, 8 turns,
+GSPO-token loss, RLOO advantages, asymmetric clipping (`0.003`/`0.005`),
+`sequence-mean-token-mean` aggregation, `low_var_kl`, temperature `0.7`, AdamW
+(`1e-6`, `b1=0.9`, `b2=0.95`, no weight decay), and gradient clipping at 100.
+Maps use the same seed/size/frozen-probability distribution as the original
+dataset recipe, but are generated directly in memory without Grain, pandas, or
+Parquet.
+
+Install the FrozenLake and distributed extras, then launch from an 8-chip TPU
+host:
+
+```bash
+pip install -e '.[frozenlake,experimental]'
+cd tunix/experimental/examples/frozenlake_dist
+WEIGHT_SYNC_MODE=raiden ./launcher.sh
+```
+
+For a small infrastructure smoke test, use a smaller supported model and turn
+off weight synchronization:
+
+```bash
+MODEL_NAME=Qwen3-1.7B MODEL_ID=Qwen/Qwen3-1.7B \
+TRAINER_TPU_CHIPS=0,1 TRAINER_TP=2 \
+ROLLOUT_TPU_CHIPS=2,3 ROLLOUT_TP=2 \
+BATCH_SIZE=1 MINI_BATCH_SIZE=1 NUM_GENERATIONS=2 \
+MAX_STEPS=1 MAX_TURNS=3 WEIGHT_SYNC_MODE=none ./launcher.sh
+```
+
+`BATCH_SIZE` is the full/global batch and determines checkpointing, global
+step advancement, and weight-sync cadence. `MINI_BATCH_SIZE` determines each
+optimizer update, so one full step performs
+`BATCH_SIZE / MINI_BATCH_SIZE` updates. Multi-step training should use
+`WEIGHT_SYNC_MODE=raiden`; `none` is intended for smoke tests, and the launcher
+rejects the protocol-only `fallback` mode through the orchestrator validation.
+
+The launcher starts one actor and one rollout worker, so `BETA` must remain
+zero. A nonzero KL coefficient requires adding a reference inference worker.
+FrozenLake is deterministic by default, matching the reference recipe's
+environment construction; set `IS_SLIPPERY=1` to enable Gymnasium's slippery
+transitions.
+
+One reference-only feature is not modeled separately: the original recipe's
+`sampler_is` threshold is folded into the distributed path's rollout-logprob
+importance ratio because Orchestrator V2 currently exposes
+`use_rollout_logps`, but not an independent sampler-IS threshold.