diff --git a/reasoning_gym/algorithmic/__init__.py b/reasoning_gym/algorithmic/__init__.py index 2cc35e61..e6188eb2 100644 --- a/reasoning_gym/algorithmic/__init__.py +++ b/reasoning_gym/algorithmic/__init__.py @@ -15,7 +15,7 @@ from .count_primes import CountPrimesConfig, CountPrimesCurriculum, CountPrimesD from .cryptarithm import CryptarithmConfig, CryptarithmDataset from .game_of_life import GameOfLifeConfig, GameOfLifeDataset from .game_of_life_halting import GameOfLifeHaltingConfig, GameOfLifeHaltingDataset -from .graph_color import GraphColorConfig, GraphColorDataset +from .graph_color import GraphColorConfig, GraphColorCurriculum, GraphColorDataset from .group_anagrams import GroupAnagramsConfig, GroupAnagramsCurriculum, GroupAnagramsDataset from .isomorphic_strings import IsomorphicStringsConfig, IsomorphicStringsCurriculum, IsomorphicStringsDataset from .jugs import JugsConfig, JugsDataset @@ -113,6 +113,7 @@ __all__ = [ "CountPrimesCurriculum", "GraphColorConfig", "GraphColorDataset", + "GraphColorCurriculum", "StringInsertionConfig", "StringInsertionDataset", "StringManipulationConfig", diff --git a/reasoning_gym/algorithmic/graph_color.py b/reasoning_gym/algorithmic/graph_color.py index fa34b63c..d1e17bf0 100644 --- a/reasoning_gym/algorithmic/graph_color.py +++ b/reasoning_gym/algorithmic/graph_color.py @@ -3,6 +3,7 @@ from dataclasses import dataclass from random import Random from typing import Any, Optional +from ..coaching import AttributeType, BaseCurriculum, RangeAttributeDefinition from ..factory import ProceduralDataset, register_dataset @@ -154,8 +155,10 @@ def greedy_graph_coloring(puzzle): class GraphColorConfig: """Configuration for GraphColor puzzle generation""" - num_colors: int = 4 - num_vertices: int = 10 + min_num_colors: int = 3 + max_num_colors: int = 3 + min_num_vertices: int = 10 + max_num_vertices: int = 10 edge_probability: float = 0.4 seed: Optional[int] = None size: int = 500 @@ -187,9 +190,9 @@ class GraphColorDataset(ProceduralDataset): while solution is None: puzzle = generate_graph_coloring_puzzle( rng=rng, - num_vertices=self.config.num_vertices, + num_vertices=rng.randint(self.config.min_num_vertices, self.config.max_num_vertices), edge_probability=self.config.edge_probability, - num_colors=self.config.num_colors, + num_colors=rng.randint(self.config.min_num_colors, self.config.max_num_colors), ) solution = greedy_graph_coloring(puzzle) @@ -237,4 +240,32 @@ Return your solution as a JSON map of vertices to colors. (For example: {{"0": 1 return 0.0 -register_dataset("graph_color", GraphColorDataset, GraphColorConfig) +class GraphColorCurriculum(BaseCurriculum): + def __init__(self): + super().__init__(GraphColorCurriculum.__name__, GraphColorConfig) + + self._define_attributes( + RangeAttributeDefinition( + name="num_vertices", + levels=[10, 20, 30, 40], + default_level=0, + description="Number of vertices in the graph", + attr_type=AttributeType.APPEND, + min_value=10, + lower_field_name="min_num_vertices", + upper_field_name="max_num_vertices", + ), + RangeAttributeDefinition( + name="num_colors", + levels=[3, 4, 5, 6], + default_level=0, + description="Number of colors in the graph", + attr_type=AttributeType.APPEND, + min_value=2, + lower_field_name="min_num_colors", + upper_field_name="max_num_colors", + ), + ) + + +register_dataset("graph_color", GraphColorDataset, GraphColorConfig, GraphColorCurriculum) diff --git a/tests/test_graph_color.py b/tests/test_graph_color.py index 01e80e04..7f631e6f 100644 --- a/tests/test_graph_color.py +++ b/tests/test_graph_color.py @@ -2,12 +2,20 @@ import json import pytest -from reasoning_gym.algorithmic.graph_color import GraphColorConfig, GraphColorDataset +from reasoning_gym.algorithmic.graph_color import GraphColorConfig, GraphColorCurriculum, GraphColorDataset def test_graph_color(): """Test basic properties and solution of generated items""" - config = GraphColorConfig(seed=42, size=10, num_vertices=10, num_colors=4, edge_probability=0.4) + config = GraphColorConfig( + seed=42, + size=10, + min_num_vertices=10, + max_num_vertices=10, + min_num_colors=4, + max_num_colors=4, + edge_probability=0.4, + ) dataset = GraphColorDataset(config) # easy @@ -22,7 +30,15 @@ def test_graph_color(): assert dataset.score_answer(answer=None, entry=item) == 0.0 # medium - config = GraphColorConfig(seed=42, size=1, num_vertices=10, num_colors=3, edge_probability=0.3) + config = GraphColorConfig( + seed=42, + size=1, + min_num_vertices=10, + max_num_vertices=10, + min_num_colors=3, + max_num_colors=3, + edge_probability=0.3, + ) dataset = GraphColorDataset(config) for item in dataset: @@ -30,7 +46,15 @@ def test_graph_color(): assert dataset.score_answer(answer=None, entry=item) == 0.0 # hard - config = GraphColorConfig(seed=42, size=1, num_vertices=40, num_colors=4, edge_probability=0.2) + config = GraphColorConfig( + seed=42, + size=1, + min_num_vertices=40, + max_num_vertices=40, + min_num_colors=4, + max_num_colors=4, + edge_probability=0.2, + ) dataset = GraphColorDataset(config) for item in dataset: @@ -38,9 +62,63 @@ def test_graph_color(): assert dataset.score_answer(answer=None, entry=item) == 0.0 # v hard - config = GraphColorConfig(seed=42, size=1, num_vertices=50, num_colors=3, edge_probability=0.1) + config = GraphColorConfig( + seed=42, + size=1, + min_num_vertices=50, + max_num_vertices=50, + min_num_colors=3, + max_num_colors=3, + edge_probability=0.1, + ) dataset = GraphColorDataset(config) for item in dataset: assert dataset.score_answer(answer=json.dumps(item["metadata"]["possible_answer"]), entry=item) == 1.0 assert dataset.score_answer(answer=None, entry=item) == 0.0 + + +def test_graph_color_curriculum(): + curriculum = GraphColorCurriculum() + + base_value = {"size": 150, "seed": 1} + + base_cfg: GraphColorConfig = curriculum.generate_configuration(base_value) + assert base_cfg.size == 150 + assert base_cfg.seed == 1 + assert base_cfg.min_num_vertices == base_cfg.max_num_vertices == 10 + assert base_cfg.min_num_colors == base_cfg.max_num_colors == 3 + + curriculum.increment_attr_level("num_vertices") + cfg = curriculum.generate_configuration(base_value) + assert cfg.min_num_vertices == 10 + assert cfg.max_num_vertices == 20 + + curriculum.increment_attr_level("num_colors") + cfg = curriculum.generate_configuration(base_value) + assert cfg.min_num_colors == 3 + assert cfg.max_num_colors == 4 + + curriculum.increment_attr_level("num_vertices") + curriculum.increment_attr_level("num_colors") + cfg = curriculum.generate_configuration(base_value) + assert cfg.min_num_vertices == 10 + assert cfg.max_num_vertices == 30 + assert cfg.min_num_colors == 3 + assert cfg.max_num_colors == 5 + + curriculum.increment_attr_level("num_vertices") + curriculum.increment_attr_level("num_colors") + cfg = curriculum.generate_configuration(base_value) + assert cfg.min_num_vertices == 10 + assert cfg.max_num_vertices == 40 + assert cfg.min_num_colors == 3 + assert cfg.max_num_colors == 6 + + curriculum.decrement_attr_level("num_vertices") + curriculum.decrement_attr_level("num_colors") + cfg = curriculum.generate_configuration(base_value) + assert cfg.min_num_vertices == 10 + assert cfg.max_num_vertices == 30 + assert cfg.min_num_colors == 3 + assert cfg.max_num_colors == 5