diff --git a/reasoning_gym/algorithmic/__init__.py b/reasoning_gym/algorithmic/__init__.py index f7f8d161..0b5c4769 100644 --- a/reasoning_gym/algorithmic/__init__.py +++ b/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", diff --git a/reasoning_gym/algorithmic/pool_matrix.py b/reasoning_gym/algorithmic/pool_matrix.py index 706d5c42..7e2acabc 100644 --- a/reasoning_gym/algorithmic/pool_matrix.py +++ b/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) diff --git a/tests/test_pool_matrix.py b/tests/test_pool_matrix.py index a529c05f..caa569d3 100644 --- a/tests/test_pool_matrix.py +++ b/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