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

fix: Ensure color cube rotation dataset always includes valid rotations

Browse source
This commit is contained in:
Andreas Koepf (aider) 2025-01-24 17:57:02 +01:00
parent 7715cd228f
commit a4181a5ad4
1 changed files with 8 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

13
reasoning_gym/cognition/color_cube_rotation.py
View file

@ -119,11 +119,14 @@ class ColorCubeRotationDataset(ProceduralDataset):
# Generate sequence of rotations
num_rotations = rng.randint(self.config.min_rotations, self.config.max_rotations)
rotations = []
for _ in range(num_rotations):
from_side = rng.choice(list(Side))
if from_side != Side.TOP: # Skip meaningless top-to-top rotation
rotations.append(from_side)
self._rotate_to_top(cube, from_side)
# Keep trying until we have at least one valid rotation
while len(rotations) < num_rotations:
# Get all sides except TOP
available_sides = [s for s in Side if s != Side.TOP]
from_side = rng.choice(available_sides)
rotations.append(from_side)
self._rotate_to_top(cube, from_side)
# Select target side for question
target_side = rng.choice(list(Side))