diff --git a/reasoning_gym/cognition/__init__.py b/reasoning_gym/cognition/__init__.py index e6cb9bce..2d3b87a6 100644 --- a/reasoning_gym/cognition/__init__.py +++ b/reasoning_gym/cognition/__init__.py @@ -6,6 +6,14 @@ Cognition tasks for training reasoning capabilities: - Working memory """ -from .number_sequences import NumberSequenceConfig, NumberSequenceDataset, sequence_dataset +from .color_cube_rotation import ColorCubeRotationConfig, ColorCubeRotationDataset, color_cube_rotation_dataset +from .number_sequences import NumberSequenceConfig, NumberSequenceDataset, number_sequence_dataset -__all__ = ["NumberSequenceDataset", "NumberSequenceConfig", "sequence_dataset"] +__all__ = [ + "NumberSequenceConfig", + "NumberSequenceDataset", + "number_sequence_dataset", + "ColorCubeRotationConfig", + "ColorCubeRotationDataset", + "color_cube_rotation_dataset", +] diff --git a/reasoning_gym/cognition/color_cube_rotation.py b/reasoning_gym/cognition/color_cube_rotation.py index 7b6ba7d5..aa58636b 100644 --- a/reasoning_gym/cognition/color_cube_rotation.py +++ b/reasoning_gym/cognition/color_cube_rotation.py @@ -27,6 +27,7 @@ class Side(StrEnum): @dataclass class Cube: """Represents a cube with colored sides""" + colors: Dict[Side, Color] def rotate_front_to_top(self) -> None: @@ -78,6 +79,7 @@ class Cube: @dataclass class ColorCubeRotationConfig: """Configuration for color cube rotation task generation""" + min_rotations: int = 1 max_rotations: int = 3 seed: Optional[int] = None @@ -99,11 +101,11 @@ class ColorCubeRotationDataset(ProceduralDataset): def __getitem__(self, idx: int) -> dict: rng = random.Random(self.seed + idx) - + # Generate initial cube state cube = self._generate_cube(rng) initial_state = cube.colors.copy() - + # Generate sequence of rotations num_rotations = rng.randint(self.config.min_rotations, self.config.max_rotations) rotations = [] @@ -112,22 +114,22 @@ class ColorCubeRotationDataset(ProceduralDataset): if from_side != Side.TOP: # Skip meaningless top-to-top rotation rotations.append(from_side) self._rotate_to_top(cube, from_side) - + # Select target side for question target_side = rng.choice(list(Side)) - + # Generate story story = self._generate_story(initial_state, rotations, target_side) - + return { "question": story, "answer": cube.colors[target_side], "metadata": { - "initial_state": {k.value: v.value for k,v in initial_state.items()}, + "initial_state": {k.value: v.value for k, v in initial_state.items()}, "rotations": [r.value for r in rotations], "target_side": target_side.value, "num_rotations": num_rotations, - } + }, } def _generate_cube(self, rng: random.Random) -> Cube: @@ -148,24 +150,23 @@ class ColorCubeRotationDataset(ProceduralDataset): if from_side in rotation_map: rotation_map[from_side]() - def _generate_story(self, initial_state: Dict[Side, Color], - rotations: List[Side], target_side: Side) -> str: + def _generate_story(self, initial_state: Dict[Side, Color], rotations: List[Side], target_side: Side) -> str: """Generate story describing cube state and rotations""" # Describe initial state story_parts = ["A cube has:"] for side in Side: story_parts.append(f"- a {initial_state[side].value} {side.value} side") - + # Describe rotations for from_side in rotations: story_parts.append( f"\nThe cube is rotated so that the side which was before at the {from_side.value} " "is now at the top." ) - + # Ask question story_parts.append(f"\nWhat is now the color of the {target_side.value} side of the cube?") - + return "\n".join(story_parts) diff --git a/reasoning_gym/cognition/number_sequences.py b/reasoning_gym/cognition/number_sequences.py index cfd80815..c674ed49 100644 --- a/reasoning_gym/cognition/number_sequences.py +++ b/reasoning_gym/cognition/number_sequences.py @@ -200,7 +200,7 @@ class NumberSequenceDataset(ProceduralDataset): } -def sequence_dataset( +def number_sequence_dataset( min_terms: int = 4, max_terms: int = 8, min_value: int = -100, diff --git a/reasoning_gym/graphs/__init__.py b/reasoning_gym/graphs/__init__.py index bf643127..399e0101 100644 --- a/reasoning_gym/graphs/__init__.py +++ b/reasoning_gym/graphs/__init__.py @@ -1,11 +1,3 @@ -from .family_relationships import ( - FamilyRelationshipsDataset, - FamilyRelationshipsConfig, - family_relationships_dataset -) +from .family_relationships import FamilyRelationshipsConfig, FamilyRelationshipsDataset, family_relationships_dataset -__all__ = [ - "FamilyRelationshipsDataset", - "FamilyRelationshipsConfig", - "family_relationships_dataset" -] +__all__ = ["FamilyRelationshipsDataset", "FamilyRelationshipsConfig", "family_relationships_dataset"] diff --git a/tests/test_color_cube_rotation.py b/tests/test_color_cube_rotation.py index fbee824e..fc62423b 100644 --- a/tests/test_color_cube_rotation.py +++ b/tests/test_color_cube_rotation.py @@ -1,31 +1,27 @@ import pytest -from reasoning_gym.cognition.color_cube_rotation import ( - color_cube_rotation_dataset, - Color, - Side, - Cube, -) + +from reasoning_gym.cognition.color_cube_rotation import Color, Cube, Side, color_cube_rotation_dataset def test_color_cube_rotation_generation(): dataset = color_cube_rotation_dataset(seed=42, size=10) - + for item in dataset: # Check required keys exist assert "question" in item assert "answer" in item assert "metadata" in item - + # Validate story format question = item["question"] assert "A cube has:" in question assert "side" in question assert "rotated" in question assert "What is now the color" in question - + # Validate answer is a valid color assert item["answer"] in [c.value for c in Color] - + # Validate metadata assert "initial_state" in item["metadata"] assert "rotations" in item["metadata"] @@ -38,7 +34,7 @@ def test_color_cube_rotation_config(): # Test invalid config raises assertion with pytest.raises(AssertionError): dataset = color_cube_rotation_dataset(min_rotations=0) - + with pytest.raises(AssertionError): dataset = color_cube_rotation_dataset(max_rotations=1, min_rotations=2) @@ -46,7 +42,7 @@ def test_color_cube_rotation_config(): def test_deterministic_generation(): dataset1 = color_cube_rotation_dataset(seed=42, size=5) dataset2 = color_cube_rotation_dataset(seed=42, size=5) - + for i in range(5): assert dataset1[i]["question"] == dataset2[i]["question"] assert dataset1[i]["answer"] == dataset2[i]["answer"] @@ -54,15 +50,17 @@ def test_deterministic_generation(): def test_cube_rotations(): # Test individual rotation operations - cube = Cube({ - Side.TOP: Color.RED, - Side.RIGHT: Color.GREEN, - Side.FRONT: Color.BLUE, - Side.LEFT: Color.YELLOW, - Side.BACK: Color.WHITE, - Side.BOTTOM: Color.ORANGE, - }) - + cube = Cube( + { + Side.TOP: Color.RED, + Side.RIGHT: Color.GREEN, + Side.FRONT: Color.BLUE, + Side.LEFT: Color.YELLOW, + Side.BACK: Color.WHITE, + Side.BOTTOM: Color.ORANGE, + } + ) + # Test front to top rotation original = cube.colors.copy() cube.rotate_front_to_top() @@ -71,4 +69,4 @@ def test_cube_rotations(): assert cube.colors[Side.BOTTOM] == original[Side.BACK] assert cube.colors[Side.BACK] == original[Side.TOP] assert cube.colors[Side.RIGHT] == original[Side.RIGHT] # Unchanged - assert cube.colors[Side.LEFT] == original[Side.LEFT] # Unchanged + assert cube.colors[Side.LEFT] == original[Side.LEFT] # Unchanged