From a4181a5ad4eecdd0508f96918a7d73617c10aee7 Mon Sep 17 00:00:00 2001 From: "Andreas Koepf (aider)" Date: Fri, 24 Jan 2025 17:57:02 +0100 Subject: [PATCH] fix: Ensure color cube rotation dataset always includes valid rotations --- reasoning_gym/cognition/color_cube_rotation.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/reasoning_gym/cognition/color_cube_rotation.py b/reasoning_gym/cognition/color_cube_rotation.py index c5bfbc71..ef2210ee 100644 --- a/reasoning_gym/cognition/color_cube_rotation.py +++ b/reasoning_gym/cognition/color_cube_rotation.py @@ -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))