diff --git a/reasoning_gym/cognition/rubiks_cube.py b/reasoning_gym/cognition/rubiks_cube.py index f39d6176..17f45315 100644 --- a/reasoning_gym/cognition/rubiks_cube.py +++ b/reasoning_gym/cognition/rubiks_cube.py @@ -116,7 +116,8 @@ class RubiksCubeDataset(ProceduralDataset): # Test the solution try: - eval_cube.rotate(answer) + expanded_answer = self.expand_moves(answer) + eval_cube.rotate(expanded_answer) solved = eval_cube.is_done() if solved: @@ -135,6 +136,22 @@ class RubiksCubeDataset(ProceduralDataset): ansi_escape = re.compile(r"(?:\x1B[@-_]|[\x80-\x9F])[0-?]*[ -/]*[@-~]") return ansi_escape.sub("", line) + def expand_moves(self, move_str): + moves = move_str.split() + expanded = [] + for move in moves: + # Split the move into the base part and any trailing digits + match = re.fullmatch(r"^([^\d]*)(\d*)$", move) + if match: + base, num_part = match.groups() + if num_part: + # Append two copies of the base if there was a number. I don't think F3 is a valid signmaster notation etc + expanded.append(base) + expanded.append(base) + else: + expanded.append(base) + return " ".join(expanded).strip() + # Register the dataset register_dataset("rubiks_cube", RubiksCubeDataset, RubiksCubeConfig)