feat(env): Binary Matrix Curriculum (#279)

* binary matrix curriculum

* register BinaryMatrixCurriculum

---------

Co-authored-by: Andreas Koepf <andreas.koepf@provisio.com>

reasoning_gym/algorithmic/__init__.py

@@ -9,7 +9,7 @@ Algorithmic tasks for training reasoning capabilities:
 from .ab import ABConfig, ABDataset
 from .base_conversion import BaseConversionConfig, BaseConversionDataset
 from .binary_alternation import BinaryAlternationConfig, BinaryAlternationCurriculum, BinaryAlternationDataset
-from .binary_matrix import BinaryMatrixConfig, BinaryMatrixDataset
+from .binary_matrix import BinaryMatrixConfig, BinaryMatrixCurriculum, BinaryMatrixDataset
 from .caesar_cipher import CaesarCipherConfig, CaesarCipherDataset
 from .count_primes import CountPrimesConfig, CountPrimesDataset
 from .cryptarithm import CryptarithmConfig, CryptarithmDataset
@@ -89,6 +89,7 @@ __all__ = [
     "ManipulateMatrixDataset",
     "BinaryMatrixConfig",
     "BinaryMatrixDataset",
+    "BinaryMatrixCurriculum",
     "PoolMatrixConfig",
     "PoolMatrixDataset",
     "ABConfig",

reasoning_gym/algorithmic/binary_matrix.py

@@ -9,6 +9,7 @@ from dataclasses import dataclass
 from random import Random
 from typing import Any, Optional
 
+from ..coaching import AttributeType, BaseCurriculum, RangeAttributeDefinition, ScalarAttributeDefinition
 from ..factory import ProceduralDataset, register_dataset
 
 QUESTION_TEMPLATE = """Given a square matrix, your job is to find the taxicab (Manhattan) distance of the nearest 0 for each cell.
@@ -44,9 +45,8 @@ class BinaryMatrixDataset(ProceduralDataset):
     def __init__(self, config: BinaryMatrixConfig):
         super().__init__(config=config, seed=config.seed, size=config.size)
 
-    def _get_binary_matrix(self, rng: Random) -> list[list[int]]:
+    def _get_binary_matrix(self, rng: Random, n: int) -> list[list[int]]:
         """Generate a random binary matrix"""
-        n = rng.randint(self.config.min_n, self.config.max_n)
         # Ensure at least one 0 in the matrix, so that a solution exists
         numbers = [0] + [0 if rng.random() < self.config.p_zero else 1 for _ in range(n**2 - 1)]
         rng.shuffle(numbers)
@@ -117,7 +117,8 @@ class BinaryMatrixDataset(ProceduralDataset):
         """Generate a single Binary Matrix question"""
         rng = Random(self.seed + idx)
 
-        matrix = self._get_binary_matrix(rng)
+        n = rng.randint(self.config.min_n, self.config.max_n)
+        matrix = self._get_binary_matrix(rng, n)
         matrix_str = self._matrix_to_str(matrix)
 
         answer = self._get_distances(matrix)
@@ -126,8 +127,42 @@ class BinaryMatrixDataset(ProceduralDataset):
         return {
             "question": QUESTION_TEMPLATE.format(matrix=matrix_str),
             "answer": answer_str,
-            "metadata": {"matrix": matrix, "solution": answer},
+            "metadata": {
+                "matrix": matrix,
+                "solution": answer,
+                "difficulty": {
+                    "n": n,
+                    "p_zero": self.config.p_zero,
+                },
+            },
         }
 
 
-register_dataset("binary_matrix", BinaryMatrixDataset, BinaryMatrixConfig)
+class BinaryMatrixCurriculum(BaseCurriculum):
+    def __init__(self):
+        super().__init__(BinaryMatrixCurriculum.__name__, BinaryMatrixConfig)
+
+        self._define_attributes(
+            ScalarAttributeDefinition(
+                name="p_zero",
+                field_name="p_zero",
+                levels=[0.5, 0.25, 0.1, 0.05],
+                default_level=0,
+                description="Board size",
+                attr_type=AttributeType.STATIC,
+                min_value=0,
+            ),
+            RangeAttributeDefinition(
+                name="n",
+                levels=[10, 50, 250, 1000],
+                default_level=0,
+                description="Board size",
+                attr_type=AttributeType.APPEND,
+                min_value=1,
+                lower_field_name="min_n",
+                upper_field_name="max_n",
+            ),
+        )
+
+
+register_dataset("binary_matrix", BinaryMatrixDataset, BinaryMatrixConfig, BinaryMatrixCurriculum)

tests/test_binary_matrix.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from reasoning_gym.algorithmic.binary_matrix import BinaryMatrixConfig, BinaryMatrixDataset
+from reasoning_gym.algorithmic.binary_matrix import BinaryMatrixConfig, BinaryMatrixCurriculum, BinaryMatrixDataset
 
 
 def test_binary_matrix_config_validation():
@@ -121,3 +121,28 @@ def test_binary_matrix_answer():
     answer = None
     entry = {"answer": "0 0 0\n0 1 0\n1 2 1"}
     assert dataset.score_answer(answer, entry) == 0.0
+
+
+def test_n_queens_curriculum():
+    curriculum = BinaryMatrixCurriculum()
+
+    base_value = {"size": 150, "seed": 1}
+
+    base_cfg: BinaryMatrixConfig = curriculum.generate_configuration(base_value)
+    assert base_cfg.seed == 1
+    assert base_cfg.size == 150
+    assert base_cfg.p_zero == 0.5
+    assert base_cfg.min_n == 10 and base_cfg.max_n == 10
+
+    # test incrementing attribute levels for n and p_zero
+    curriculum.increment_attr_level("n")
+    curriculum.increment_attr_level("p_zero")
+    increased_cfg = curriculum.generate_configuration(base_value)
+    assert increased_cfg.p_zero == 0.25
+    assert increased_cfg.min_n == 10 and increased_cfg.max_n == 50
+
+    # test decrementing attribute level for n again
+    curriculum.decrement_attr_level("n")
+    partially_decreased_cfg = curriculum.generate_configuration(base_value)
+    assert partially_decreased_cfg.p_zero == 0.25
+    assert partially_decreased_cfg.min_n == 10 and partially_decreased_cfg.max_n == 10