add survo env (#461)

* add survo env

* add survo curriculum

* add survo tests

Author: olliestanley

reasoning_gym/games/__init__.py

@@ -20,6 +20,7 @@
 from .rush_hour import RushHourConfig, RushHourCurriculum, RushHourDataset
 from .sokoban import SokobanConfig, SokobanCurriculum, SokobanDataset
 from .sudoku import SudokuConfig, SudokuCurriculum, SudokuDataset
+from .survo import SurvoConfig, SurvoCurriculum, SurvoDataset
 from .tower_of_hanoi import HanoiConfig, HanoiCurriculum, HanoiDataset
 from .tsumego import TsumegoConfig, TsumegoCurriculum, TsumegoDataset
 
@@ -45,12 +46,15 @@
     "Puzzle24Config",
     "Puzzle24Dataset",
     "Puzzle24Curriculum",
-    "SudokuConfig",
-    "SudokuCurriculum",
-    "SudokuDataset",
     "SokobanConfig",
     "SokobanCurriculum",
     "SokobanDataset",
+    "SudokuConfig",
+    "SudokuCurriculum",
+    "SudokuDataset",
+    "SurvoConfig",
+    "SurvoCurriculum",
+    "SurvoDataset",
     "RushHourConfig",
     "RushHourCurriculum",
     "RushHourDataset",

reasoning_gym/games/survo.py

@@ -0,0 +1,186 @@
+"""
+Survo dataset, adapted for Reasoning Gym from SynthRL: https://github.com/MiniMax-AI/SynLogic/tree/main/games/tasks/survo
+"""
+
+from dataclasses import dataclass
+from random import Random
+from typing import Any, Optional
+
+import numpy as np
+
+from ..coaching import BaseCurriculum, RangeAttributeDefinition
+from ..factory import ProceduralDataset, register_dataset
+
+DATASET_NAME = "survo"
+
+PROMPT_TEMPLATES = [
+    "Given a {n}*{n} matrix where the last element of each row and column equals the sum of the other elements in that row or column. The matrix is:\n{matrix}\nwhere some elements are replaced with 0. You have a set of numbers {numbers} that can be filled into the 0 positions to satisfy the rules. Please fill in the matrix. Each number can only be used once.",
+    "You have a {n}*{n} matrix with some positions already filled with numbers and others marked with 0. The matrix is:\n{matrix}\nThe last number in each row and column represents the sum of all other numbers in that row or column. You need to fill in the 0 positions using the numbers {numbers} to satisfy these conditions. Each number can only be used once.",
+    "Complete the following Survo puzzle. In this {n}*{n} matrix:\n{matrix}\nthe cells marked with 0 need to be filled with numbers. The last number in each row and column equals the sum of all other numbers in that row or column. You can use the following numbers: {numbers}. Each number can only be used once.",
+    "In this {n}*{n} Survo matrix puzzle:\n{matrix}\nthe 0 cells need to be filled with numbers from the set {numbers}. The last element in each row and column is the sum of all other elements in that row or column. Each number can only be used once. Provide the completed matrix.",
+    "Solve this {n}*{n} matrix puzzle:\n{matrix}\nwhere 0 represents empty cells that need to be filled. The last number in each row and column equals the sum of all other numbers in that row or column. You have the numbers {numbers} to place in the empty cells. Each number can only be used once.",
+]
+
+
+@dataclass
+class SurvoConfig:
+    min_board_size: int = 4
+    max_board_size: int = 5
+    min_empty: int = 3
+    max_empty: int = 5
+    min_num: int = 1
+    max_num: int = 9
+    seed: Optional[int] = None
+    size: int = 500  # Virtual dataset size
+
+    def validate(self):
+        """Validate configuration parameters"""
+        assert self.min_board_size > 3, "min_board_size must be greater than 3"
+        assert self.max_board_size >= self.min_board_size, "max_board_size must be >= min_board_size"
+        assert self.min_empty > 0, "min_empty must be > 0"
+        assert self.max_empty <= (self.min_board_size - 1) * (
+            self.min_board_size - 1
+        ), f"max_empty must be <= {(self.min_board_size - 1) * (self.min_board_size - 1)}"
+        assert self.min_empty <= self.max_empty, "min_empty must be <= max_empty"
+        assert self.min_num > 0, "min_num must be > 0"
+        assert self.min_num < self.max_num, "min_num must be less than max_num"
+
+
+class SurvoDataset(ProceduralDataset):
+    def __init__(self, config: SurvoConfig):
+        super().__init__(config=config, seed=config.seed, size=config.size)
+
+    def __len__(self) -> int:
+        return self.config.size
+
+    def __iter__(self):
+        self._current_idx = 0
+        return self
+
+    def __next__(self):
+        if self._current_idx >= self.config.size:
+            raise StopIteration
+        item = self[self._current_idx]
+        self._current_idx += 1
+        return item
+
+    def __getitem__(self, idx: int) -> dict:
+        rng = Random(self.config.seed + idx)
+
+        board_size = rng.randint(self.config.min_board_size, self.config.max_board_size)
+        num_empty = rng.randint(self.config.min_empty, self.config.max_empty)
+
+        filled_matrix, puzzle, candidate_numbers = self._generate_valid_matrix(
+            rng, board_size, num_empty, self.config.min_num, self.config.max_num
+        )
+
+        puzzle_str = "\n".join(" ".join(str(x) for x in row) for row in puzzle)
+        solution_str = "\n".join(" ".join(str(x) for x in row) for row in filled_matrix)
+
+        question = rng.choice(PROMPT_TEMPLATES).format(n=board_size, matrix=puzzle_str, numbers=candidate_numbers)
+
+        return {
+            "question": question,
+            "answer": solution_str,
+            "metadata": {
+                "source_dataset": DATASET_NAME,
+                "source_idx": idx,
+                "puzzle": puzzle.tolist(),
+                "solution": filled_matrix.tolist(),
+                "candidate_numbers": candidate_numbers,
+                "board_size": board_size,
+                "num_empty": num_empty,
+                "min_num": self.config.min_num,
+                "max_num": self.config.max_num,
+                "difficulty": {
+                    "board_size": (self.config.min_board_size, self.config.max_board_size),
+                    "empty": (self.config.min_empty, self.config.max_empty),
+                },
+            },
+        }
+
+    def _generate_valid_matrix(
+        self, rng: Random, n: int, num_empty: int, min_num: int, max_num: int
+    ) -> tuple[np.ndarray, np.ndarray, list[int]]:
+        matrix = np.zeros((n, n), dtype=int)
+
+        for i in range(n - 1):
+            for j in range(n - 1):
+                matrix[i, j] = rng.randint(min_num, max_num)
+
+        for i in range(n - 1):
+            row_sum = sum(matrix[i, 0 : n - 1])
+            matrix[i, n - 1] = row_sum
+
+            col_sum = sum(matrix[0 : n - 1, i])
+            matrix[n - 1, i] = col_sum
+
+        matrix[n - 1, n - 1] = sum(matrix[0 : n - 1, n - 1])
+
+        filled_matrix = matrix.copy()
+
+        positions = [(i, j) for i in range(n - 1) for j in range(n - 1)]
+        selected_positions = rng.sample(positions, num_empty)
+
+        candidate_numbers = []
+        for i, j in selected_positions:
+            candidate_numbers.append(int(matrix[i, j]))
+            matrix[i, j] = 0
+
+        return filled_matrix, matrix, candidate_numbers
+
+    def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float:
+        if not isinstance(answer, str):
+            return 0.0
+
+        board_size = entry["metadata"]["board_size"]
+        grid = self._parse_grid(answer)
+        true_grid = entry["metadata"]["solution"]
+
+        if len(grid) != board_size or any(len(row) != board_size for row in grid):
+            return 0.0
+
+        for i in range(board_size):
+            for j in range(board_size):
+                if grid[i][j] != true_grid[i][j]:
+                    return 0.0
+
+        return 1.0
+
+    def _parse_grid(self, answer: str) -> list[list[str]]:
+        grid = []
+        for line in answer.strip().split("\n"):
+            row = []
+            for c in line.strip().split():
+                try:
+                    row.append(int(c))
+                except ValueError:
+                    continue  # Ignore non-integer values
+            if row:
+                grid.append(row)
+        return grid
+
+
+class SurvoCurriculum(BaseCurriculum):
+    def __init__(self):
+        super().__init__(SurvoCurriculum.__name__, SurvoConfig)
+
+        self._define_attributes(
+            RangeAttributeDefinition(
+                name="board_size",
+                levels=[4, 5, 6, 7],
+                description="Board size (n x n)",
+                lower_field_name="min_board_size",
+                upper_field_name="max_board_size",
+            ),
+            RangeAttributeDefinition(
+                name="empty",
+                levels=[4, 9, 16, 25],
+                description="Number of empty cells",
+                lower_field_name="min_empty",
+                upper_field_name="max_empty",
+            ),
+        )
+
+
+register_dataset(DATASET_NAME, SurvoDataset, SurvoConfig, SurvoCurriculum)

tests/test_survo.py

@@ -0,0 +1,151 @@
+import numpy as np
+import pytest
+
+from reasoning_gym.coaching.base_curriculum import DefaultCurriculumContext, RangeAttributeMode
+from reasoning_gym.games import SurvoConfig, SurvoCurriculum, SurvoDataset
+
+
+def test_survo_config_validation():
+    """Bad configs should raise."""
+    # min_board_size must be > 3
+    with pytest.raises(AssertionError):
+        SurvoConfig(min_board_size=3).validate()
+
+    # max_board_size must be ≥ min_board_size
+    with pytest.raises(AssertionError):
+        SurvoConfig(min_board_size=6, max_board_size=5).validate()
+
+    # min_empty ≤ max_empty and within board-area limits
+    with pytest.raises(AssertionError):
+        SurvoConfig(min_empty=6, max_empty=5).validate()
+
+    # min_num < max_num
+    with pytest.raises(AssertionError):
+        SurvoConfig(min_num=5, max_num=5).validate()
+
+
+def test_survo_deterministic():
+    """Same seed ⇒ identical items."""
+    cfg = SurvoConfig(seed=123, size=15, min_board_size=4, max_board_size=5, min_empty=3, max_empty=5)
+    ds1, ds2 = SurvoDataset(cfg), SurvoDataset(cfg)
+
+    for i in range(len(ds1)):
+        assert ds1[i] == ds2[i]
+
+
+def test_survo_items():
+    """Generated items have expected structure and metadata."""
+    cfg = SurvoConfig(seed=99, size=20, min_board_size=4, max_board_size=5, min_empty=3, max_empty=5)
+    ds = SurvoDataset(cfg)
+
+    for itm in ds:
+        md = itm["metadata"]
+
+        # Basic keys
+        assert set(itm.keys()) == {"question", "answer", "metadata"}
+        assert "puzzle" in md and "solution" in md and "candidate_numbers" in md
+
+        orig = np.array(md["puzzle"])
+        full = np.array(md["solution"])
+
+        # Dimensions
+        n = full.shape[0]
+        assert cfg.min_board_size <= n <= cfg.max_board_size
+        assert orig.shape == full.shape == (n, n)
+
+        # Number of empties
+        empties = np.count_nonzero(orig == 0)
+        assert empties == md["num_empty"]
+        assert cfg.min_empty <= empties <= cfg.max_empty
+
+        # Candidate numbers should match removed values (order disregarded)
+        removed_values = full[orig == 0].tolist()
+        assert sorted(removed_values) == sorted(md["candidate_numbers"])
+
+
+def test_survo_solution_validity():
+    """Solution must satisfy Survo row/column-sum rules."""
+    cfg = SurvoConfig(seed=321, size=10, min_board_size=4, max_board_size=5, min_empty=3, max_empty=5)
+    ds = SurvoDataset(cfg)
+
+    for itm in ds:
+        m = np.array(itm["metadata"]["solution"])
+        n = m.shape[0]
+
+        # Row sums (exclude last row)
+        for r in range(n - 1):
+            assert m[r, : n - 1].sum() == m[r, n - 1]
+
+        # Column sums (exclude last col)
+        for c in range(n - 1):
+            assert m[: n - 1, c].sum() == m[n - 1, c]
+
+        # Grand total cell
+        assert m[n - 1, n - 1] == m[: n - 1, n - 1].sum()
+
+
+def test_survo_difficulty_levels():
+    """More allowed empties ⇒ puzzles have, on average, more blank cells."""
+    seed, n_items, board = 777, 8, 5
+
+    def avg_empties(min_empty, max_empty):
+        cfg = SurvoConfig(
+            seed=seed,
+            size=n_items,
+            min_board_size=board,
+            max_board_size=board,
+            min_empty=min_empty,
+            max_empty=max_empty,
+        )
+        ds = SurvoDataset(cfg)
+        return np.mean([np.count_nonzero(np.array(itm["metadata"]["puzzle"]) == 0) for itm in ds])
+
+    low = avg_empties(3, 3)
+    mid = avg_empties(6, 6)
+    high = avg_empties(10, 10)
+
+    assert low < mid < high
+
+
+def test_survo_answer_scoring():
+    """Correct answer ⇒ 1.0; variations score lower."""
+    cfg = SurvoConfig(seed=42, size=5, min_board_size=4, max_board_size=4, min_empty=3, max_empty=3)
+    ds = SurvoDataset(cfg)
+
+    for itm in ds:
+        correct = itm["answer"]
+        assert ds.score_answer(correct, itm) == 1.0
+
+        # Tamper with a single cell
+        candidate_numbers = itm["metadata"]["candidate_numbers"]
+        wrong = correct.replace(str(candidate_numbers[0]), str(max(candidate_numbers) + 1), 1)
+        assert ds.score_answer(wrong, itm) < 1.0
+
+        # Bad type / empty
+        assert ds.score_answer(None, itm) == 0.0
+        assert ds.score_answer("", itm) == 0.0
+
+
+def test_survo_curriculum():
+    """SurvoCurriculum controls board size and empties as advertised."""
+    cur = SurvoCurriculum()
+    base_val = {"size": 100, "seed": 1}
+    ctx = DefaultCurriculumContext(mode=RangeAttributeMode.UPPER_BOUND)
+
+    # Level 0
+    cfg0: SurvoConfig = cur.generate_configuration(base_val, context=ctx)
+    assert cfg0.min_board_size == cfg0.max_board_size == 4
+    assert cfg0.min_empty == cfg0.max_empty == 4
+
+    # Increment levels
+    cur.increment_attr_level("board_size")
+    cur.increment_attr_level("empty")
+    cfg1: SurvoConfig = cur.generate_configuration(base_val, context=ctx)
+    assert cfg1.min_board_size == cfg1.max_board_size == 5
+    assert cfg1.min_empty == cfg1.max_empty == 9
+
+    # Global progression
+    cur.set_global_level(3)
+    cfg_max: SurvoConfig = cur.generate_configuration(base_val, context=ctx)
+    assert cfg_max.min_board_size == cfg_max.max_board_size == 7
+    assert cfg_max.min_empty == cfg_max.max_empty == 25