open-thought/reasoning-gym
1
0
Fork 0
mirror of https://github.com/open-thought/reasoning-gym.git synced 2026-04-28 17:29:39 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

color cube rotation curriculum (#347)

Browse source
This commit is contained in:
Zafir Stojanovski 2025-03-13 21:04:34 +01:00 • committed by GitHub
parent ec3e414a8c
commit 9fcc277101
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 53 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

3
reasoning_gym/cognition/__init__.py
View file

@ -2,7 +2,7 @@
Cognition tasks for training reasoning capabilities.
"""
from .color_cube_rotation import ColorCubeRotationConfig, ColorCubeRotationDataset
from .color_cube_rotation import ColorCubeRotationConfig, ColorCubeRotationCurriculum, ColorCubeRotationDataset
from .figlet_fonts import FigletFontConfig, FigletFontDataset
from .modulo_grid import ModuloGridConfig, ModuloGridDataset
from .needle_haystack import NeedleHaystackConfig, NeedleHaystackDataset
@ -13,6 +13,7 @@ from .rubiks_cube import RubiksCubeConfig, RubiksCubeDataset
__all__ = [
"ColorCubeRotationConfig",
"ColorCubeRotationDataset",
"ColorCubeRotationCurriculum",
"FigletFontConfig",
"FigletFontDataset",
"NumberSequenceConfig",

23
reasoning_gym/cognition/color_cube_rotation.py
View file

@ -3,6 +3,7 @@ from dataclasses import dataclass
from enum import StrEnum
from typing import Any, Optional
from ..coaching import AttributeType, BaseCurriculum, RangeAttributeDefinition
from ..factory import ProceduralDataset, register_dataset
@ -140,6 +141,7 @@ class ColorCubeRotationDataset(ProceduralDataset):
"rotations": [r.value for r in rotations],
"target_side": target_side.value,
"num_rotations": num_rotations,
"difficulty": {"rotations": num_rotations},
},
}
@ -204,4 +206,23 @@ class ColorCubeRotationDataset(ProceduralDataset):
return reward
register_dataset("color_cube_rotation", ColorCubeRotationDataset, ColorCubeRotationConfig)
class ColorCubeRotationCurriculum(BaseCurriculum):
def __init__(self):
super().__init__(ColorCubeRotationCurriculum.__name__, ColorCubeRotationConfig)
# Define attributes
self._define_attributes(
RangeAttributeDefinition(
name="rotations",
levels=[1, 5, 10, 50, 100],
default_level=1,
description="Number of rotations to perform on the cube",
attr_type=AttributeType.APPEND,
min_value=1,
lower_field_name="min_rotations",
upper_field_name="max_rotations",
)
)
register_dataset("color_cube_rotation", ColorCubeRotationDataset, ColorCubeRotationConfig, ColorCubeRotationCurriculum)