feat(env): NQueens Curriculum (#262)

* curriculum & tests
This commit is contained in:
Zafir Stojanovski 2025-03-05 15:05:17 +01:00 committed by GitHub
parent 5d7fbac0ad
commit 8ecc723607
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 59 additions and 3 deletions

View file

@ -1,10 +1,11 @@
from .attributes import AttributeDefinition, AttributeType, RangeAttributeDefinition
from .attributes import AttributeDefinition, AttributeType, RangeAttributeDefinition, ScalarAttributeDefinition
from .base_curriculum import BaseCurriculum
from .coach import Coach, GroupedScores, ScoreBoard, ScoreStats
__all__ = [
"AttributeType",
"AttributeDefinition",
"ScalarAttributeDefinition",
"RangeAttributeDefinition",
"BaseCurriculum",
"Coach",

View file

@ -13,7 +13,7 @@ from .knight_swap import KnightSwapConfig, KnightSwapDataset
from .mahjong import MahjongPuzzleConfig, MahjongPuzzleDataset
from .maze import MazeConfig, MazeDataset
from .mini_sudoku import MiniSudokuConfig, MiniSudokuDataset
from .n_queens import NQueensDataset
from .n_queens import NQueensConfig, NQueensCurriculum, NQueensDataset
from .rush_hour import RushHourConfig, RushHourDataset
from .sokoban import SokobanConfig, SokobanDataset
from .sudoku import SudokuConfig, SudokuDataset
@ -40,6 +40,8 @@ __all__ = [
"HanoiConfig",
"HanoiDataset",
"NQueensDataset",
"NQueensConfig",
"NQueensCurriculum",
"TsumegoConfig",
"TsumegoDataset",
"KnightSwapConfig",

View file

@ -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
MIN_BOARD_SIZE = 4
@ -151,4 +152,31 @@ class NQueensDataset(ProceduralDataset):
return 0.0
class NQueensCurriculum(BaseCurriculum):
def __init__(self):
super().__init__(NQueensCurriculum.__name__, NQueensConfig)
self._define_attributes(
ScalarAttributeDefinition(
name="n",
field_name="n",
levels=[4, 6, 8, 12],
default_level=0,
description="Board size",
attr_type=AttributeType.STATIC,
min_value=4,
),
RangeAttributeDefinition(
name="num_removed",
levels=[2, 4, 6, 10],
default_level=0,
description="Number of queens to remove",
attr_type=AttributeType.APPEND,
min_value=1,
lower_field_name="min_remove",
upper_field_name="max_remove",
),
)
register_dataset("n_queens", NQueensDataset, NQueensConfig)

View file

@ -2,7 +2,7 @@
import pytest
from reasoning_gym.games.n_queens import NQueensConfig, NQueensDataset
from reasoning_gym.games.n_queens import NQueensConfig, NQueensCurriculum, NQueensDataset
def test_nqueens_config_validation():
@ -146,3 +146,28 @@ def is_valid_solution(board: list[list[str]]) -> bool:
off_diags.add(r - c)
return num_queens == n
def test_n_queens_curriculum():
curriculum = NQueensCurriculum()
base_value = {"size": 150, "seed": 1}
base_cfg: NQueensConfig = curriculum.generate_configuration(base_value)
assert base_cfg.seed == 1
assert base_cfg.size == 150
assert base_cfg.n == 4
assert base_cfg.min_remove == 2 and base_cfg.max_remove == 2
# test incrementing attribute levels for n & num_removed attributes
curriculum.increment_attr_level("n")
curriculum.increment_attr_level("num_removed")
increased_cfg = curriculum.generate_configuration(base_value)
assert increased_cfg.n == 6
assert increased_cfg.min_remove == 2 and increased_cfg.max_remove == 4
# test decrementing attribute level for n again
curriculum.decrement_attr_level("n")
partially_decreased_cfg = curriculum.generate_configuration(base_value)
assert partially_decreased_cfg.n == 4
assert partially_decreased_cfg.min_remove == 2 and partially_decreased_cfg.max_remove == 4