diff --git a/reasoning_gym/algorithmic/__init__.py b/reasoning_gym/algorithmic/__init__.py index a62aed26..b54511f1 100644 --- a/reasoning_gym/algorithmic/__init__.py +++ b/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", diff --git a/reasoning_gym/algorithmic/binary_matrix.py b/reasoning_gym/algorithmic/binary_matrix.py index 772cde97..9a1c2e32 100644 --- a/reasoning_gym/algorithmic/binary_matrix.py +++ b/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) diff --git a/tests/test_binary_matrix.py b/tests/test_binary_matrix.py index 88ef9adc..7d07ce09 100644 --- a/tests/test_binary_matrix.py +++ b/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