pool matrix curriculum (#298)

reasoning_gym/algorithmic/__init__.py

@@ -26,7 +26,7 @@ from .number_filtering import NumberFilteringConfig, NumberFilteringDataset
 from .number_sorting import NumberSortingConfig, NumberSortingDataset
 from .palindrome_generation import PalindromeConfig, PalindromeDataset
 from .palindrome_partitioning import PalindromePartitioningConfig, PalindromePartitioningDataset
-from .pool_matrix import PoolMatrixConfig, PoolMatrixDataset
+from .pool_matrix import PoolMatrixConfig, PoolMatrixCurriculum, PoolMatrixDataset
 from .ransom_note import RansomNoteConfig, RansomNoteDataset
 from .rotate_matrix import RotateMatrixConfig, RotateMatrixCurriculum, RotateMatrixDataset
 from .rotten_oranges import RottenOrangesConfig, RottenOrangesCurriculum, RottenOrangesDataset
@@ -99,6 +99,7 @@ __all__ = [
     "BinaryMatrixCurriculum",
     "PoolMatrixConfig",
     "PoolMatrixDataset",
+    "PoolMatrixCurriculum",
     "ABConfig",
     "ABDataset",
     "CountPrimesConfig",

reasoning_gym/algorithmic/pool_matrix.py

@@ -6,6 +6,7 @@ from typing import Any, Optional
 
 import numpy as np
 
+from ..coaching import AttributeType, BaseCurriculum, RangeAttributeDefinition
 from ..factory import ProceduralDataset, register_dataset
 
 QUESTION_TEMPLATE = """Your job is to perform max/average pooling on the given matrix.
@@ -25,9 +26,10 @@ class PoolMatrixConfig:
     """Configuration for Pool Matrix dataset generation"""
 
     min_rows: int = 2  # Minimum rows of the matrix
-    min_cols: int = 2  # Minimum columns of the matrix
     max_rows: int = 10  # Maximum rows of the matrix
+    min_cols: int = 2  # Minimum columns of the matrix
     max_cols: int = 10  # Maximum columns of the matrix
+    min_pool_size: int = 1  # Minimum pooling size
     max_pool_size: int = 3  # Maximum pooling size
 
     size: int = 500  # Virtual dataset size
@@ -36,10 +38,11 @@ class PoolMatrixConfig:
     def validate(self):
         """Validate configuration parameters"""
         assert 2 <= self.min_rows, "min_rows must be at least 2"
-        assert 2 <= self.min_cols, "min_cols must be at least 2"
         assert self.min_rows <= self.max_rows, "max_rows must be at least min_rows"
+        assert 2 <= self.min_cols, "min_cols must be at least 2"
         assert self.min_cols <= self.max_cols, "max_cols must be at least min_cols"
-        assert 1 <= self.max_pool_size, "max_pool_size must be at least 1"
+        assert 1 <= self.min_pool_size, "min_pool_size must be at least 1"
+        assert self.min_pool_size <= self.max_pool_size, "max_pool_size must be at least min_pool_size"
 
 
 class PoolMatrixDataset(ProceduralDataset):
@@ -48,12 +51,6 @@ class PoolMatrixDataset(ProceduralDataset):
     def __init__(self, config: PoolMatrixConfig):
         super().__init__(config=config, seed=config.seed, size=config.size)
 
-    def _get_matrix(self, rng: Random) -> np.ndarray:
-        """Generate a random matrix"""
-        rows = rng.randint(self.config.min_rows, self.config.max_rows)
-        cols = rng.randint(self.config.min_rows, self.config.max_cols)
-        return np.random.randint(0, 10, (rows, cols))
-
     def _matrix_to_str(self, matrix: np.ndarray) -> str:
         """Get a string representation of the matrix"""
         return "\n".join(" ".join(str(round(x, 2)) for x in row) for row in matrix)
@@ -101,10 +98,12 @@ class PoolMatrixDataset(ProceduralDataset):
         rng = Random(self.seed + idx)
         np.random.seed(self.seed + idx)
 
-        matrix = self._get_matrix(rng)
+        rows = rng.randint(self.config.min_rows, self.config.max_rows)
+        cols = rng.randint(self.config.min_rows, self.config.max_cols)
+        matrix = np.random.randint(0, 10, (rows, cols))
         matrix_str = self._matrix_to_str(matrix)
 
-        pool_size = rng.randint(1, self.config.max_pool_size)
+        pool_size = rng.randint(self.config.min_pool_size, self.config.max_pool_size)
         pool_type = rng.choice(["average", "max"])
 
         answer = self._average_pool(matrix, pool_size) if pool_type == "average" else self._max_pool(matrix, pool_size)
@@ -118,8 +117,51 @@ class PoolMatrixDataset(ProceduralDataset):
                 "pool_type": pool_type,
                 "pool_size": pool_size,
                 "solution": answer.tolist(),
+                "difficulty": {
+                    "rows": rows,
+                    "cols": cols,
+                    "pool_size": pool_size,
+                },
             },
         }
 
 
-register_dataset("pool_matrix", PoolMatrixDataset, PoolMatrixConfig)
+class PoolMatrixCurriculum(BaseCurriculum):
+    def __init__(self):
+        super().__init__(PoolMatrixCurriculum.__name__, PoolMatrixConfig)
+
+        self._define_attributes(
+            RangeAttributeDefinition(
+                name="rows",
+                levels=[10, 25, 50, 100],
+                default_level=0,
+                description="Board size",
+                attr_type=AttributeType.APPEND,
+                min_value=2,
+                lower_field_name="min_rows",
+                upper_field_name="max_rows",
+            ),
+            RangeAttributeDefinition(
+                name="cols",
+                levels=[10, 25, 50, 100],
+                default_level=0,
+                description="Board size",
+                attr_type=AttributeType.APPEND,
+                min_value=2,
+                lower_field_name="min_cols",
+                upper_field_name="max_cols",
+            ),
+            RangeAttributeDefinition(
+                name="pool_size",
+                levels=[3, 5, 7, 9],
+                default_level=0,
+                description="Pool size",
+                attr_type=AttributeType.APPEND,
+                min_value=1,
+                lower_field_name="min_pool_size",
+                upper_field_name="max_pool_size",
+            ),
+        )
+
+
+register_dataset("pool_matrix", PoolMatrixDataset, PoolMatrixConfig, PoolMatrixCurriculum)

tests/test_pool_matrix.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 
-from reasoning_gym.algorithmic.pool_matrix import PoolMatrixConfig, PoolMatrixDataset
+from reasoning_gym.algorithmic.pool_matrix import PoolMatrixConfig, PoolMatrixCurriculum, PoolMatrixDataset
 
 
 def test_pool_matrix_config_validation():
@@ -161,3 +161,32 @@ def test_pool_matrix_int_answer():
                 matrix = matrix.reshape(1, 1)
             int_answer = "\n".join(" ".join(str(x) for x in row) for row in matrix)
             assert dataset.score_answer(answer=int_answer, entry=entry) == 1.0
+
+
+def test_pool_matrix_curriculum():
+    curriculum = PoolMatrixCurriculum()
+
+    base_value = {"size": 150, "seed": 1}
+
+    base_cfg: PoolMatrixConfig = curriculum.generate_configuration(base_value)
+    assert base_cfg.seed == 1
+    assert base_cfg.size == 150
+    assert base_cfg.min_rows == 10 and base_cfg.max_rows == 10
+    assert base_cfg.min_cols == 10 and base_cfg.max_cols == 10
+    assert base_cfg.min_pool_size == 3 and base_cfg.max_pool_size == 3
+
+    # test incrementing attribute levels
+    curriculum.increment_attr_level("rows")
+    curriculum.increment_attr_level("cols")
+    curriculum.increment_attr_level("pool_size")
+    increased_cfg = curriculum.generate_configuration(base_value)
+    assert increased_cfg.min_rows == 10 and increased_cfg.max_rows == 25
+    assert increased_cfg.min_cols == 10 and increased_cfg.max_cols == 25
+    assert increased_cfg.min_pool_size == 3 and increased_cfg.max_pool_size == 5
+
+    # test decrementing attribute level for pool_size again
+    curriculum.decrement_attr_level("pool_size")
+    partially_decreased_cfg = curriculum.generate_configuration(base_value)
+    assert partially_decreased_cfg.min_rows == 10 and partially_decreased_cfg.max_rows == 25
+    assert partially_decreased_cfg.min_cols == 10 and partially_decreased_cfg.max_cols == 25
+    assert partially_decreased_cfg.min_pool_size == 3 and partially_decreased_cfg.max_pool_size == 3