diff --git a/reasoning_gym/graphs/__init__.py b/reasoning_gym/graphs/__init__.py index d8e2b825..1018e33a 100644 --- a/reasoning_gym/graphs/__init__.py +++ b/reasoning_gym/graphs/__init__.py @@ -1,6 +1,7 @@ from .course_schedule import CourseScheduleConfig, CourseScheduleCurriculum, CourseScheduleDataset from .family_relationships import FamilyRelationshipsConfig, FamilyRelationshipsCurriculum, FamilyRelationshipsDataset from .largest_island import LargestIslandConfig, LargestIslandCurriculum, LargestIslandDataset +from .path_star import PathStarConfig, PathStarCurriculum, PathStarDataset from .quantum_lock import QuantumLockConfig, QuantumLockCurriculum, QuantumLockDataset from .shortest_path import ShortestPathConfig, ShortestPathCurriculum, ShortestPathDataset @@ -14,6 +15,9 @@ __all__ = [ "LargestIslandDataset", "LargestIslandConfig", "LargestIslandCurriculum", + "PathStarConfig", + "PathStarDataset", + "PathStarCurriculum", "CourseScheduleDataset", "CourseScheduleConfig", "CourseScheduleCurriculum", diff --git a/reasoning_gym/graphs/path_star.py b/reasoning_gym/graphs/path_star.py new file mode 100644 index 00000000..53be8845 --- /dev/null +++ b/reasoning_gym/graphs/path_star.py @@ -0,0 +1,161 @@ +""" +Pathfinding problems in a path-star graph structure. +Inspired by https://arxiv.org/pdf/2403.06963 +""" + +import random +from dataclasses import dataclass +from typing import Any, Optional + +from ..coaching import BaseCurriculum, RangeAttributeDefinition, ScalarAttributeDefinition +from ..factory import ProceduralDataset, register_dataset + +DATASET_NAME = "path_star" + +PROMPT_TEMPLATE = """ +Find a path from the start node to the goal node in the following path-star graph. +Respond with only the sequence of node labels in the path, including the start and goal nodes. +Separate node labels with a single space. + +The graph is represented as a list of edges, where each edge is defined by two node labels. +The edges are separated by a vertical bar '|'. Then, the start and goal nodes are specified after a slash '/'. + +Example: +|1 2|1 3|2 4|3 5/1 5 = 1 3 5 + +Solve the following task: +{task} +""" + + +@dataclass +class PathStarConfig: + """Configuration for Path Star dataset generation""" + + degree: int = 3 + node_range: int = 100_000 + min_path_length: int = 3 + max_path_length: int = 5 + + reversed: bool = False + + size: int = 500 # Virtual dataset size + seed: Optional[int] = None + + def validate(self) -> None: + """Validate configuration parameters""" + assert self.degree >= 2, "degree must be at least 2" + assert self.min_path_length >= 2, "min_path_length must be at least 2" + assert self.min_path_length <= self.max_path_length, "min_path_length must be <= max_path_length" + assert ( + self.node_range > self.degree * self.max_path_length + 1 + ), "node_range must exceed degree * max_path_length + 1 for unique labels" + + +class PathStarDataset(ProceduralDataset): + """Procedurally generates path-star graph problems.""" + + def __init__(self, config: PathStarConfig): + super().__init__(config=config, seed=config.seed, size=config.size) + + def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float: + """Score an answer. Path is unique in a star graph, so only exact match counts.""" + if not isinstance(answer, str) or len(answer.strip()) == 0: + return 0.0 + # Normalize: strip, collapse whitespace + answer_normalized = " ".join(answer.strip().split()) + oracle_normalized = " ".join(entry["answer"].strip().split()) + if answer_normalized == oracle_normalized: + return 1.0 + return 0.0 + + def __getitem__(self, idx: int) -> dict[str, Any]: + rng = random.Random(self.seed + idx) + + cfg: PathStarConfig = self.config + center = rng.randrange(cfg.node_range) + path_length = rng.randint(cfg.min_path_length, cfg.max_path_length) + + # allocate unique node labels + paths = [] + used = {center} + for _ in range(cfg.degree): + path = [] + for _ in range(path_length): + n = rng.randrange(cfg.node_range) + while n in used: + n = rng.randrange(cfg.node_range) + used.add(n) + path.append(n) + paths.append(path) + + goal_path = rng.choice(paths) + goal = goal_path[-1] + + # build edge list + edges = [(center, p[0]) for p in paths] + for p in paths: + edges.extend(zip(p[:-1], p[1:])) + rng.shuffle(edges) + + edges_str = "".join(f"|{u} {v}" for u, v in edges) + if cfg.reversed: + prefix = f"{edges_str}/{goal} {center} = " + else: + prefix = f"{edges_str}/{center} {goal} = " + question = PROMPT_TEMPLATE.format(task=prefix) + + # gold path + gold = [center] + goal_path + if cfg.reversed: + gold = list(reversed(gold)) + answer = " ".join(map(str, gold)) + + return { + "question": question, + "answer": answer, + "metadata": { + "source_dataset": DATASET_NAME, + "source_index": idx, + "center": center, + "goal": goal, + "path_length": path_length, + "goal_path": gold if not cfg.reversed else list(reversed(gold)), + "difficulty": { + "degree": cfg.degree, + "node_range": cfg.node_range, + "path_length": (cfg.min_path_length, cfg.max_path_length), + }, + }, + } + + +class PathStarCurriculum(BaseCurriculum): + def __init__(self): + super().__init__(PathStarCurriculum.__name__, PathStarConfig) + + # Define attributes + self._define_attributes( + ScalarAttributeDefinition( + name="degree", + levels=[2, 3, 4, 5], + description="Degree of the graph", + field_name="degree", + ), + ScalarAttributeDefinition( + name="node_range", + levels=[10_000, 50_000, 100_000, 200_000], + description="Range of node labels", + field_name="node_range", + ), + RangeAttributeDefinition( + name="path_length", + levels=[3, 5, 6, 7], + description="Length of paths in the graph", + lower_field_name="min_path_length", + upper_field_name="max_path_length", + ), + ) + + +register_dataset(DATASET_NAME, PathStarDataset, PathStarConfig, PathStarCurriculum) diff --git a/tests/test_path_star.py b/tests/test_path_star.py new file mode 100644 index 00000000..5436813a --- /dev/null +++ b/tests/test_path_star.py @@ -0,0 +1,195 @@ +"""Tests for Path Star graph problem generation""" + +import pytest + +from reasoning_gym.graphs.path_star import PathStarConfig, PathStarCurriculum, PathStarDataset + + +def test_path_star_config_validation(): + """Test that invalid configs raise appropriate errors""" + with pytest.raises(AssertionError): + config = PathStarConfig(degree=1) # Must be >= 2 + config.validate() + + with pytest.raises(AssertionError): + config = PathStarConfig(min_path_length=1) # Must be >= 2 + config.validate() + + with pytest.raises(AssertionError): + config = PathStarConfig(min_path_length=5, max_path_length=3) # min > max + config.validate() + + with pytest.raises(AssertionError): + config = PathStarConfig(degree=3, max_path_length=5, node_range=16) # node_range too small (need > 3*5+1=16) + config.validate() + + +def test_path_star_dataset_deterministic(): + """Test that dataset generates same items with same seed""" + config = PathStarConfig(seed=42, size=10) + dataset1 = PathStarDataset(config) + dataset2 = PathStarDataset(config) + + for i in range(len(dataset1)): + assert dataset1[i] == dataset2[i] + + +def test_path_star_dataset_items(): + """Test basic properties of generated items""" + config = PathStarConfig(min_path_length=3, max_path_length=5, size=10, seed=42) + dataset = PathStarDataset(config) + + for i in range(len(dataset)): + item = dataset[i] + # Check item structure + assert isinstance(item, dict) + assert "question" in item + assert "answer" in item + assert "metadata" in item + + # Check metadata fields + assert item["metadata"]["source_dataset"] == "path_star" + assert item["metadata"]["source_index"] == i + assert "center" in item["metadata"] + assert "goal" in item["metadata"] + assert "path_length" in item["metadata"] + assert "goal_path" in item["metadata"] + assert "difficulty" in item["metadata"] + + # Verify answer format: space-separated integers + answer_parts = item["answer"].split() + assert all(part.isdigit() for part in answer_parts) + + # First node should be center, last should be goal + center = item["metadata"]["center"] + goal = item["metadata"]["goal"] + assert int(answer_parts[0]) == center + assert int(answer_parts[-1]) == goal + + # Path length should match: center + path_length nodes + path_length = item["metadata"]["path_length"] + assert len(answer_parts) == path_length + 1 + + # Path length within configured range + assert config.min_path_length <= path_length <= config.max_path_length + + +def test_path_star_dataset_iteration(): + """Test that iteration respects dataset size""" + config = PathStarConfig(size=5, seed=42) + dataset = PathStarDataset(config) + + items = list(dataset) + assert len(items) == config.size + + # Test multiple iterations yield same items + assert items == list(dataset) + + +def test_path_star_answer_correctness(): + """Test that generated paths are valid by checking edge connectivity""" + config = PathStarConfig(size=20, seed=123) + dataset = PathStarDataset(config) + + for i in range(len(dataset)): + item = dataset[i] + question = item["question"] + answer_parts = [int(x) for x in item["answer"].split()] + + # Parse edges from the question + # Format: ...edges_str/start goal = ... + # Extract the task part between "Solve the following task:\n" and the end + task_line = question.split("Solve the following task:\n")[1].strip() + edge_part, _ = task_line.split("/") + edges = set() + for edge_str in edge_part.split("|"): + edge_str = edge_str.strip() + if edge_str: + u, v = edge_str.split() + edges.add((int(u), int(v))) + + # Verify consecutive nodes in the answer are connected by edges + for j in range(len(answer_parts) - 1): + u, v = answer_parts[j], answer_parts[j + 1] + assert (u, v) in edges, f"Edge ({u}, {v}) not found in edges for item {i}" + + +def test_path_star_score_answer(): + """Test the score_answer method""" + config = PathStarConfig(seed=42, size=5) + dataset = PathStarDataset(config) + item = dataset[0] + oracle = item["answer"] + + # Exact match + assert dataset.score_answer(oracle, item) == 1.0 + + # Match with extra whitespace + assert dataset.score_answer(f" {oracle} ", item) == 1.0 + + # Match with extra internal whitespace + spaced = oracle.replace(" ", " ") + assert dataset.score_answer(spaced, item) == 1.0 + + # Wrong answer + assert dataset.score_answer("0 1 2 3", item) == 0.0 + + # None + assert dataset.score_answer(None, item) == 0.0 + + # Empty string + assert dataset.score_answer("", item) == 0.0 + + +def test_path_star_reversed(): + """Test that reversed=True produces correct answer and task format""" + config_fwd = PathStarConfig(seed=42, size=5, reversed=False) + config_rev = PathStarConfig(seed=42, size=5, reversed=True) + dataset_fwd = PathStarDataset(config_fwd) + dataset_rev = PathStarDataset(config_rev) + + for i in range(len(dataset_fwd)): + item_fwd = dataset_fwd[i] + item_rev = dataset_rev[i] + + # Reversed answer should be the forward answer reversed + fwd_parts = item_fwd["answer"].split() + rev_parts = item_rev["answer"].split() + assert rev_parts == list(reversed(fwd_parts)) + + # Task format should swap start/goal + center = item_fwd["metadata"]["center"] + goal = item_fwd["metadata"]["goal"] + assert f"/{center} {goal} = " in item_fwd["question"] + assert f"/{goal} {center} = " in item_rev["question"] + + +def test_path_star_curriculum(): + """Test curriculum creates valid configs at various levels""" + curriculum = PathStarCurriculum() + + base_value = {"size": 150, "seed": 1} + + # Level 0 (base) + base_cfg: PathStarConfig = curriculum.generate_configuration(base_value) + assert base_cfg.seed == 1 + assert base_cfg.size == 150 + assert base_cfg.degree == 2 + assert base_cfg.node_range == 10_000 + assert base_cfg.min_path_length == 3 and base_cfg.max_path_length == 3 + + # Increment attributes + curriculum.increment_attr_level("degree") + curriculum.increment_attr_level("node_range") + curriculum.increment_attr_level("path_length") + increased_cfg = curriculum.generate_configuration(base_value) + assert increased_cfg.degree == 3 + assert increased_cfg.node_range == 50_000 + assert increased_cfg.min_path_length == 3 and increased_cfg.max_path_length == 5 + + # Decrement degree back + curriculum.decrement_attr_level("degree") + partial_cfg = curriculum.generate_configuration(base_value) + assert partial_cfg.degree == 2 + assert partial_cfg.node_range == 50_000 + assert partial_cfg.min_path_length == 3 and partial_cfg.max_path_length == 5