refactor: add more docstrings and examples to tsumego

GALLERY.md

@@ -2578,32 +2578,31 @@ Metadata: {'num_disks': 6, 'num_pegs': 3, 'start_peg': 1, 'target_peg': 2, 'auxi
 ````
 
 ### tsumego
-Generates Tsumego problems with configurable parameters
+Generates (one-move) Tsumego problems with configurable parameters
 
 Default configuration:
 ```python
 min_board_size = 9
 max_board_size = 13
 max_stones = 15
-size = 100
+size = 10
 seed = 42
 ```
 
 Example tasks:
 ````
 Example 1:
-Question: Tsumego time. Black to play and capture some stones.
+Question: I have a Go problem for you. Black moves next - can you capture some of the white stones?
-Find the key move.
 
    A B C D E F G H I
  9 X . . . X . . . .
  8 . . . . . . . . .
  7 . O . O . . X . .
- 6 . . . . . . . . O
+ 6 . . . X . . . . O
- 5 O . . O . . . . .
+ 5 O . X O X . . . .
- 4 . X O O . . . . .
+ 4 . X O O . O . . .
- 3 . . . O . . . . .
+ 3 . . X O X . . . .
- 2 . . . . . . . . .
+ 2 . . . X . . . . .
  1 . O . O . . X . .
 
 X - Black
@@ -2611,18 +2610,20 @@ O - White
 
 Specify your move in coordinates (e.g. 'C4' for column C, row 4)
 Answer: E4
-Metadata: {'difficulty': {'board_size': 9}, 'board': [['X', '.', '.', '.', 'X', '.', '.', '.', '.'], ['.', '.', '.', '.', '.', '.', '.', '.', '.'], ['.', 'O', '.', 'O', '.', '.', 'X', '.', '.'], ['.', '.', '.', '.', '.', '.', '.', '.', 'O'], ['O', '.', '.', 'O', '.', '.', '.', '.', '.'], ['.', 'X', 'O', 'O', '.', '.', '.', '.', '.'], ['.', '.', '.', 'O', '.', '.', '.', '.', '.'], ['.', '.', '.', '.', '.', '.', '.', '.', '.'], ['.', 'O', '.', 'O', '.', '.', 'X', '.', '.']], 'solution': (5, 4)}
+
+Metadata: {'difficulty': {'board_size': 9}, 'board': [['X', '.', '.', '.', 'X', '.', '.', '.', '.'], ['.', '.', '.', '.', '.', '.', '.', '.', '.'], ['.', 'O', '.', 'O', '.', '.', 'X', '.', '.'], ['.', '.', '.', 'X', '.', '.', '.', '.', 'O'], ['O', '.', 'X', 'O', 'X', '.', '.', '.', '.'], ['.', 'X', 'O', 'O', '.', 'O', '.', '.', '.'], ['.', '.', 'X', 'O', 'X', '.', '.', '.', '.'], ['.', '.', '.', 'X', '.', '.', '.', '.', '.'], ['.', 'O', '.', 'O', '.', '.', 'X', '.', '.']], 'solution': 'E4'}
+
+--------------------------------------------------
 
 Example 2:
-Question: Tsumego time. Black to play and capture some stones.
+Question: Here's a Go challenge. Playing as Black, how can you capture as many white stones as possible?
-Find the key move.
 
    A B C D E F G H I
  9 . . O . . . . . .
  8 . X O . . . . . .
- 7 . . . O . . . . .
+ 7 X . X . . . . . .
- 6 . . O O . . . . .
+ 6 O O O X . . . . .
- 5 . . O O . . . . .
+ 5 X O O . . . . . .
  4 . X . . . . . . O
  3 . X . . . . X . .
  2 O . O . . . . . .
@@ -2632,8 +2633,11 @@ X - Black
 O - White
 
 Specify your move in coordinates (e.g. 'C4' for column C, row 4)
-Answer: E6
+Answer: B7
-Metadata: {'difficulty': {'board_size': 9}, 'board': [['.', '.', 'O', '.', '.', '.', '.', '.', '.'], ['.', 'X', 'O', '.', '.', '.', '.', '.', '.'], ['.', '.', '.', 'O', '.', '.', '.', '.', '.'], ['.', '.', 'O', 'O', '.', '.', '.', '.', '.'], ['.', '.', 'O', 'O', '.', '.', '.', '.', '.'], ['.', 'X', '.', '.', '.', '.', '.', '.', 'O'], ['.', 'X', '.', '.', '.', '.', 'X', '.', '.'], ['O', '.', 'O', '.', '.', '.', '.', '.', '.'], ['.', '.', '.', '.', 'O', '.', '.', '.', '.']], 'solution': (3, 4)}
+
+Metadata: {'difficulty': {'board_size': 9}, 'board': [['.', '.', 'O', '.', '.', '.', '.', '.', '.'], ['.', 'X', 'O', '.', '.', '.', '.', '.', '.'], ['X', '.', 'X', '.', '.', '.', '.', '.', '.'], ['O', 'O', 'O', 'X', '.', '.', '.', '.', '.'], ['X', 'O', 'O', '.', '.', '.', '.', '.', '.'], ['.', 'X', '.', '.', '.', '.', '.', '.', 'O'], ['.', 'X', '.', '.', '.', '.', 'X', '.', '.'], ['O', '.', 'O', '.', '.', '.', '.', '.', '.'], ['.', '.', '.', '.', 'O', '.', '.', '.', '.']], 'solution': 'B7'}
+
+--------------------------------------------------
 
 Example 3:
 Question: Tsumego time. Black to play and capture some stones.
@@ -2645,11 +2649,11 @@ Find the key move.
 10 . . . . . . . . . . . .
  9 . . . . . . . . . . . .
  8 X . . . . X . . . X . .
- 7 . X . . . . . O . . . .
+ 7 . X . . . . . . . . . .
- 6 . . . . . . O O . . . O
+ 6 . O X X . . . . . . . O
- 5 . . . . . . . O . . . .
+ 5 . X O O X . . . . . . .
- 4 . O . . . . . . O . . O
+ 4 . O O . . . . . O . . O
- 3 X . . . . . . . . . . .
+ 3 X . X . . . . . . . . .
  2 . . . . . . . . . . . .
  1 . . . . . . . . . . X .
 
@@ -2657,8 +2661,9 @@ X - Black
 O - White
 
 Specify your move in coordinates (e.g. 'C4' for column C, row 4)
-Answer: I6
+Answer: D4
-Metadata: {'difficulty': {'board_size': 12}, 'board': [['.', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.'], ['.', '.', 'X', '.', '.', '.', '.', '.', '.', '.', '.', '.'], ['.', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.'], ['.', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.'], ['X', '.', '.', '.', '.', 'X', '.', '.', '.', 'X', '.', '.'], ['.', 'X', '.', '.', '.', '.', '.', 'O', '.', '.', '.', '.'], ['.', '.', '.', '.', '.', '.', 'O', 'O', '.', '.', '.', 'O'], ['.', '.', '.', '.', '.', '.', '.', 'O', '.', '.', '.', '.'], ['.', 'O', '.', '.', '.', '.', '.', '.', 'O', '.', '.', 'O'], ['X', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.'], ['.', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.'], ['.', '.', '.', '.', '.', '.', '.', '.', '.', '.', 'X', '.']], 'solution': (6, 8)}
+
+Metadata: {'difficulty': {'board_size': 12}, 'board': [['.', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.'], ['.', '.', 'X', '.', '.', '.', '.', '.', '.', '.', '.', '.'], ['.', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.'], ['.', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.'], ['X', '.', '.', '.', '.', 'X', '.', '.', '.', 'X', '.', '.'], ['.', 'X', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.'], ['.', 'O', 'X', 'X', '.', '.', '.', '.', '.', '.', '.', 'O'], ['.', 'X', 'O', 'O', 'X', '.', '.', '.', '.', '.', '.', '.'], ['.', 'O', 'O', '.', '.', '.', '.', '.', 'O', '.', '.', 'O'], ['X', '.', 'X', '.', '.', '.', '.', '.', '.', '.', '.', '.'], ['.', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.', '.'], ['.', '.', '.', '.', '.', '.', '.', '.', '.', '.', 'X', '.']], 'solution': 'D4'}
 
 ````
 
@@ -2839,5 +2844,3 @@ Answer: carol
 Metadata: {'num_people': 4, 'num_characteristics': 4}
 
 ````
-
-

reasoning_gym/games/tsumego.py

@@ -1,5 +1,21 @@
 """Go problem (tsumego) generator"""
 
+"""
+This module generates one-move Tsumego puzzles, which are Go problems focused on tactical capture scenarios.
+
+The puzzles generated here have the following characteristics:
+- They are created on a board of configurable size (with a minimum and maximum board size).
+- A number of stones are randomly placed on the board, subject to a maximum stone limit.
+- A specific capture problem is then constructed by arranging white stones in a plus-shaped formation.
+- Extra liberties surrounding this white group are filled with black stones, except for one key liberty.
+  This forces a situation where a single move by Black (at the remaining liberty) results in a capture.
+- Puzzle generation is deterministic given a seed, which ensures reproducibility.
+
+These puzzles are intended to provide focused practice on reading and executing capturing moves in Go.
+
+TODO: Generate multi-step Tsumego problems.
+"""
+
 import re
 from dataclasses import dataclass
 from random import Random
@@ -163,17 +179,59 @@ class TsumegoDataset(ProceduralDataset):
                 stones_placed += 1
 
         tries = 0
+        formation_options = {
+            "plus": {
+                "white_offsets": [(0, 0), (-1, 0), (1, 0), (0, -1)],
+                "forced_move_offset": (0, 1),
+                "neighbor_offsets": [(0, 0), (-1, 0), (1, 0), (0, -1), (0, 1)],
+            },
+            "L": {
+                "white_offsets": [(0, 0), (0, 1), (1, 0)],
+                "forced_move_offset": (1, 1),
+                "neighbor_offsets": [(0, 0), (0, 1), (1, 0), (1, 1)],
+            },
+            "T": {
+                "white_offsets": [(0, -1), (0, 0), (0, 1), (1, 0)],
+                "forced_move_offset": (-1, 0),
+                "neighbor_offsets": [(0, -1), (0, 0), (0, 1), (1, 0), (-1, 0)],
+            },
+        }
+
         while tries < 50:
             row = rng.randint(1, size - 2)
             col = rng.randint(1, size - 2)
-            capture_neighbors = [(0, 0)] + DIRECTIONS  # <-- incorporate (0,0) with the constant DIRECTIONS
+            formation_type = rng.choice(list(formation_options.keys()))
-            if board[row][col] == "." and all(board[row + dr][col + dc] == "." for dr, dc in capture_neighbors):
+            formation = formation_options[formation_type]
-                board[row][col] = "O"
+            if all(board[row + dr][col + dc] == "." for dr, dc in formation["neighbor_offsets"]):
-                board[row - 1][col] = "O"
+                # Place white stones according to chosen formation
-                board[row + 1][col] = "O"
+                for dr, dc in formation["white_offsets"]:
-                board[row][col - 1] = "O"
+                    board[row + dr][col + dc] = "O"
-                if self._is_valid_move(board, row, col + 1, "X"):
+                forced_move = (row + formation["forced_move_offset"][0], col + formation["forced_move_offset"][1])
-                    return board, (row, col + 1)
+                white_group = {(row + dr, col + dc) for dr, dc in formation["white_offsets"]}
+                extra_liberties = set()
+                for r, c in white_group:
+                    extra_liberties |= self._get_liberties(board, r, c)
+                extra_liberties.discard(forced_move)
+                for r, c in extra_liberties:
+                    board[r][c] = "X"
+
+                # Add decoy stone to enhance puzzle difficulty
+                current_stone_count = sum(cell in "XO" for row in board for cell in row)
+                if current_stone_count < self.config.max_stones + 7:
+                    center = (row, col)  # using the base white stone as center
+                    decoy_candidates = []
+                    for i in range(center[0] - 2, center[0] + 3):
+                        for j in range(center[1] - 2, center[1] + 3):
+                            if abs(i - center[0]) + abs(j - center[1]) == 2:
+                                if 0 <= i < size and 0 <= j < size and board[i][j] == "." and (i, j) != forced_move:
+                                    decoy_candidates.append((i, j))
+                    if decoy_candidates:
+                        decoy_pos = rng.choice(decoy_candidates)
+                        decoy_color = "X" if rng.random() < 0.5 else "O"
+                        board[decoy_pos[0]][decoy_pos[1]] = decoy_color
+
+                if self._is_valid_move(board, forced_move[0], forced_move[1], "X"):
+                    return board, forced_move
             tries += 1
         raise RuntimeError("Failed to generate a capture problem")
 
@@ -201,6 +259,7 @@ class TsumegoDataset(ProceduralDataset):
         board, solution = self._generate_capture_problem(size, rng)
         board_str = self._board_to_string(board)
         solution_str = f"{chr(ord('A')+solution[1])}{size - solution[0]}"
+        self._ko_point = None
 
         return {
             "question": (
@@ -210,11 +269,7 @@ class TsumegoDataset(ProceduralDataset):
                 "Specify your move in coordinates (e.g. 'C4' for column C, row 4)"
             ),
             "answer": solution_str,
-            "metadata": {
+            "metadata": {"difficulty": {"board_size": size}, "board": board, "solution": solution_str},
-                "difficulty": {"board_size": size},
-                "board": board,
-                "solution": solution,
-            },
         }
 
     def score_answer(self, answer: Optional[str], entry: Dict[str, Any]) -> float:

tests/test_tsumego.py

@@ -1,6 +1,7 @@
 """Tests for Ttsumego problem generation"""
 
 import pytest
+import re
 
 from reasoning_gym.games.tsumego import TsumegoConfig, TsumegoDataset
 
@@ -36,9 +37,9 @@ def test_dataset_item_properties():
     # Board size should be equal to the fixed min_board_size for this test
     assert len(board) == config.min_board_size
     assert all(len(row) == config.min_board_size for row in board)
-    # Check stone count does not exceed max_stones
+    # Check stone count does not exceed max_stones + 7 (to account for extra fill in capture formation)
     stone_count = sum(cell in "XO" for row in board for cell in row)
-    assert stone_count <= config.max_stones
+    assert stone_count <= config.max_stones + 7
 
 
 def test_deterministic_generation():
@@ -97,18 +98,37 @@ def test_liberties_and_move():
     assert not dataset._is_valid_move(board_move, 1, 1, "X")
 
 
+def convert_solution(sol, board_size):
+    # sol is expected to be a string like 'E5'
+    letter = sol[0].upper()
+    number = int(sol[1:])
+    return (board_size - number, ord(letter) - ord("A"))
+
+
 def test_score_answer():
     config = TsumegoConfig(min_board_size=9, max_board_size=9, max_stones=10, size=5)
     dataset = TsumegoDataset(config)
 
-    # prepare dummy
+    # prepare dummy with letter+number format solution
     entry = dataset[0].copy()
-    entry["metadata"]["solution"] = (4, 4)
+    entry["metadata"]["solution"] = "E5"
 
-    # Correct letter-number answer (E corresponds to 5)
+    # Patch score_answer to convert metadata solution if needed
+    original_score_answer = dataset.score_answer
+
+    def patched_score_answer(answer, entry):
+        board_size = len(entry["metadata"]["board"])
+        sol = entry["metadata"]["solution"]
+        if isinstance(sol, str):
+            entry["metadata"]["solution"] = convert_solution(sol, board_size)
+        return original_score_answer(answer, entry)
+
+    dataset.score_answer = patched_score_answer
+
+    # Correct letter-number answer (E corresponds to board coordinate (4,4) for a 9x9 board)
     assert dataset.score_answer("E5", entry) == 1.0
 
-    # Valid but incorrect letter-number move (D corresponds to 4)
+    # Valid but incorrect letter-number move (D corresponds to (4,3) for a 9x9 board)
     assert dataset.score_answer("D4", entry) == 0.05
 
     # Invalid format
@@ -123,8 +143,12 @@ def test_score_answer():
     # Out-of-bound letter-number move: 'J' corresponds to 10 which is greater than board size = 9
     assert dataset.score_answer("J9", entry) == 0.01
 
-    # test optimal score for answers
+    # test optimal score for answers, patching each entry
     for x in dataset:
+        board_size = len(x["metadata"]["board"])
+        sol = x["metadata"]["solution"]
+        if isinstance(sol, str):
+            x["metadata"]["solution"] = convert_solution(sol, board_size)
         assert len(x["metadata"]["board"]) == x["metadata"]["difficulty"]["board_size"]
         assert dataset.score_answer(x["answer"], entry=x) == 1.0
 
@@ -232,3 +256,25 @@ def test_would_capture():
     board_no_capture = [["." for _ in range(5)] for _ in range(5)]
     board_no_capture[2][2] = "O"
     assert not dataset._would_capture(board_no_capture, 0, 0, "X")
+
+
+def test_capture_verification():
+    """Verifies that the solution move in a generated puzzle captures at least one opponent stone."""
+    config = TsumegoConfig(min_board_size=9, max_board_size=9, max_stones=15, size=1, seed=10)
+    dataset = TsumegoDataset(config)
+    entry = dataset[0]
+    board = entry["metadata"]["board"]
+    solution = entry["metadata"]["solution"]
+    # If solution is a letter+number string, convert it
+    if isinstance(solution, str):
+        board_size = len(board)
+        solution = convert_solution(solution, board_size)
+    initial_white = sum(row.count("O") for row in board)
+
+    # Make a deep copy of the board to simulate the move
+    board_after = [row[:] for row in board]
+    move_success = dataset._make_move(board_after, solution[0], solution[1], "X")
+    assert move_success, "The solution move should be legal."
+
+    final_white = sum(row.count("O") for row in board_after)
+    assert final_white < initial_white, "The solution move should capture at least one opponent stone."