diff --git a/reasoning_gym/cognition/__init__.py b/reasoning_gym/cognition/__init__.py index 2839a5be..79185467 100644 --- a/reasoning_gym/cognition/__init__.py +++ b/reasoning_gym/cognition/__init__.py @@ -6,8 +6,8 @@ from .color_cube_rotation import ColorCubeRotationConfig, ColorCubeRotationDatas from .figlet_fonts import FigletFontConfig, FigletFontDataset from .modulo_grid import ModuloGridConfig, ModuloGridDataset from .needle_haystack import NeedleHaystackConfig, NeedleHaystackDataset -from .number_sequences import NumberSequenceConfig, NumberSequenceDataset -from .rectangle_count import RectangleCountConfig, RectangleCountDataset +from .number_sequences import NumberSequenceConfig, NumberSequenceCurriculum, NumberSequenceDataset +from .rectangle_count import RectangleCountConfig, RectangleCountCurriculum, RectangleCountDataset from .rubiks_cube import RubiksCubeConfig, RubiksCubeDataset __all__ = [ @@ -17,9 +17,11 @@ __all__ = [ "FigletFontDataset", "NumberSequenceConfig", "NumberSequenceDataset", + "NumberSequenceCurriculum", "RubiksCubeConfig", "RubiksCubeDataset", "RectangleCountConfig", + "RectangleCountCurriculum", "RectangleCountDataset", "NeedleHaystackConfig", "NeedleHaystackDataset", diff --git a/reasoning_gym/cognition/number_sequences.py b/reasoning_gym/cognition/number_sequences.py index 554989aa..d8ea262b 100644 --- a/reasoning_gym/cognition/number_sequences.py +++ b/reasoning_gym/cognition/number_sequences.py @@ -3,6 +3,7 @@ from enum import StrEnum from random import Random from typing import Optional +from ..coaching import AttributeType, BaseCurriculum, ScalarAttributeDefinition from ..factory import ProceduralDataset, register_dataset @@ -198,4 +199,21 @@ class NumberSequenceDataset(ProceduralDataset): } -register_dataset("number_sequence", NumberSequenceDataset, NumberSequenceConfig) +class NumberSequenceCurriculum(BaseCurriculum): + def __init__(self): + super().__init__(NumberSequenceCurriculum.__name__, NumberSequenceConfig) + + self._define_attributes( + ScalarAttributeDefinition( + name="max_complexity", + levels=[1, 2, 3, 4], + default_level=0, + description="Maximum number of operations to combine", + attr_type=AttributeType.STATIC, + min_value=1, + field_name="max_complexity", + ), + ) + + +register_dataset("number_sequence", NumberSequenceDataset, NumberSequenceConfig, NumberSequenceCurriculum) diff --git a/reasoning_gym/cognition/rectangle_count.py b/reasoning_gym/cognition/rectangle_count.py index 6f096e94..2318aa47 100644 --- a/reasoning_gym/cognition/rectangle_count.py +++ b/reasoning_gym/cognition/rectangle_count.py @@ -2,6 +2,7 @@ from dataclasses import dataclass from random import Random from typing import Any, Optional +from ..coaching import AttributeType, BaseCurriculum, ScalarAttributeDefinition from ..factory import ProceduralDataset, register_dataset QUESTION_TEMPLATE = """Your task is to count how many rectangles are present in an ASCII grid. @@ -116,7 +117,7 @@ class RectangleCountDataset(ProceduralDataset): return { "question": QUESTION_TEMPLATE.format(puzzle=puzzle), "answer": str(answer), - "metadata": {"puzzle": puzzle, "solution": answer}, + "metadata": {"puzzle": puzzle, "solution": answer, "difficulty": {"max_rectangles": target}}, } def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float: @@ -138,4 +139,22 @@ class RectangleCountDataset(ProceduralDataset): return 0.0 -register_dataset("rectangle_count", RectangleCountDataset, RectangleCountConfig) +class RectangleCountCurriculum(BaseCurriculum): + def __init__(self): + super().__init__(RectangleCountCurriculum.__name__, RectangleCountConfig) + + # Define attributes + self._define_attributes( + ScalarAttributeDefinition( + name="max_rectangles", + levels=[1, 3, 5, 10], + default_level=0, + description="Number of rectangles in the grid", + attr_type=AttributeType.STATIC, + min_value=1, + field_name="max_rectangles", + ), + ) + + +register_dataset("rectangle_count", RectangleCountDataset, RectangleCountConfig, RectangleCountCurriculum) diff --git a/tests/test_number_sequences.py b/tests/test_number_sequences.py index 69afcff3..e4b189a5 100644 --- a/tests/test_number_sequences.py +++ b/tests/test_number_sequences.py @@ -1,6 +1,12 @@ import pytest -from reasoning_gym.cognition.number_sequences import NumberSequenceConfig, NumberSequenceDataset, Operation, PatternRule +from reasoning_gym.cognition.number_sequences import ( + NumberSequenceConfig, + NumberSequenceCurriculum, + NumberSequenceDataset, + Operation, + PatternRule, +) def test_sequence_config_validation(): @@ -75,3 +81,36 @@ def test_sequence_dataset_iteration(): # Test multiple iterations yield same items assert items == list(dataset) + + +def test_number_sequence_curriculum(): + """Test the number sequence curriculum functionality""" + curriculum = NumberSequenceCurriculum() + + # Test with custom base values + base_value = {"size": 150, "seed": 42} + + # Test basic configuration generation + base_cfg: NumberSequenceConfig = curriculum.generate_configuration(base_value) + assert base_cfg.seed == 42 + assert base_cfg.size == 150 + assert base_cfg.max_complexity == 1 # Default level (0) corresponds to complexity 1 + + # Test attribute level increment + curriculum.increment_attr_level("max_complexity") + increased_cfg = curriculum.generate_configuration(base_value) + assert increased_cfg.max_complexity == 2 # Level 1 corresponds to complexity 2 + + # Test attribute level increment again + curriculum.increment_attr_level("max_complexity") + increased_cfg = curriculum.generate_configuration(base_value) + assert increased_cfg.max_complexity == 3 # Level 2 corresponds to complexity 3 + + # Test that other parameters remain unchanged + assert increased_cfg.seed == 42 + assert increased_cfg.size == 150 + + # Test attribute level decrement + curriculum.decrement_attr_level("max_complexity") + decreased_cfg = curriculum.generate_configuration(base_value) + assert decreased_cfg.max_complexity == 2 # Back to level 1, complexity 2 diff --git a/tests/test_rectangle_count.py b/tests/test_rectangle_count.py index fcbb09f0..5c7b539c 100644 --- a/tests/test_rectangle_count.py +++ b/tests/test_rectangle_count.py @@ -1,6 +1,10 @@ import pytest -from reasoning_gym.cognition.rectangle_count import RectangleCountConfig, RectangleCountDataset +from reasoning_gym.cognition.rectangle_count import ( + RectangleCountConfig, + RectangleCountCurriculum, + RectangleCountDataset, +) def test_dice(): @@ -17,3 +21,31 @@ def test_dice(): # Test the scoring assert dataset.score_answer(answer=item["answer"], entry=item) == 1.0 assert dataset.score_answer(answer=None, entry=item) == 0.0 + + +def test_rc_curriculum(): + """Test the RectangleCountCurriculum functionality""" + curriculum = RectangleCountCurriculum() + + base_value = {"size": 150, "seed": 1} + + # Test the initial configuration + base_cfg: RectangleCountConfig = curriculum.generate_configuration(base_value) + assert base_cfg.seed == 1 + assert base_cfg.size == 150 + assert base_cfg.max_rectangles == 1 # Based on number_rectangles attribute level 0 (value 1) + + # Test incrementing the number_rectangles attribute + curriculum.increment_attr_level("max_rectangles") + increased_cfg = curriculum.generate_configuration(base_value) + assert increased_cfg.max_rectangles == 3 # Based on number_rectangles attribute level 1 (value 3) + + # Test another increment + curriculum.increment_attr_level("max_rectangles") + increased_cfg = curriculum.generate_configuration(base_value) + assert increased_cfg.max_rectangles == 5 # Based on number_rectangles attribute level 2 (value 5) + + # Test decrementing + curriculum.decrement_attr_level("max_rectangles") + decreased_cfg = curriculum.generate_configuration(base_value) + assert decreased_cfg.max_rectangles == 3 # Back to level 1 (value 3)