Merge pull request #166 from joesharratt1229/fix/envs

Fix/envs

reasoning_gym/algorithmic/caesar_cipher.py

@@ -73,7 +73,7 @@ class CaesarCipherDataset(ProceduralDataset):
         cipher_text = self._caesar_encrypt(sentence, rotation)
 
         return {
-            "question": f"Decrypt this Caesar cipher text: {cipher_text}",
+            "question": f"Decrypt this Caesar cipher text: {cipher_text}. Provide only the decrypted text as your final answer.",
             "answer": sentence,
             "metadata": {"rotation": rotation, "cipher_text": cipher_text, "clear_text": sentence},
         }

reasoning_gym/cognition/color_cube_rotation.py

@@ -185,8 +185,28 @@ class ColorCubeRotationDataset(ProceduralDataset):
 
         # Ask question
         story_parts.append(f"\nWhat is now the color of the {target_side.value} side of the cube?")
+        story_parts.append(f"Provide only the color as your final answer.")
 
         return "\n".join(story_parts)
 
+    def score_answer(self, answer: Optional[str], entry: Dict[str, any]) -> float:
+        reward = 0.0
+        metadata = entry["metadata"]
+        if answer is not None:
+            try:
+                answer_formatted = answer.lower()
+                solved = answer_formatted == metadata["answer"]
+                if solved:
+                    reward = 1.0
+                elif metadata["answer"] in answer_formatted:
+                    reward = 0.25
+                elif len(answer.strip()) > 0:
+                    reward = 0.05
+                else:
+                    reward = 0.01
+            except:
+                reward = 0.01
+        return reward
+
 
 register_dataset("color_cube_rotation", ColorCubeRotationDataset, ColorCubeRotationConfig)

reasoning_gym/geometry/advanced_geometry.py

@@ -1,7 +1,9 @@
 import random
+import re
 from dataclasses import dataclass, field
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
 
+import numpy as np
 import sympy
 from sympy.geometry import Point, Segment, Triangle
 
@@ -35,6 +37,18 @@ class AdvancedGeometryConfig:
         assert len(self.task_types) > 0, "Must specify at least one task type."
 
 
+# Join format instructions into a single string
+GEOMETRY_FORMAT_INSTRUCTIONS = "\n".join(
+    [
+        "For all geometry problems:",
+        "1. Give coordinates in the form (x, y)",
+        "2. Round decimal answers to 3 decimal places",
+        "3. Use the degree symbol ° for angles",
+        "4. Return only th angle, coordinates, or radius as your answer.",
+    ]
+)
+
+
 class AdvancedGeometryDataset(ProceduralDataset):
     """
     A dataset for advanced geometry tasks using coordinate geometry.
@@ -43,16 +57,16 @@ class AdvancedGeometryDataset(ProceduralDataset):
     def __init__(self, config: AdvancedGeometryConfig):
         self._prompt_templates = {
             "orthocenter": [
-                "Given triangle ABC with coordinates A={A}, B={B}, and C={C}, find the coordinates of its orthocenter.",
-                "For triangle with vertices A={A}, B={B}, and C={C}, determine the orthocenter (intersection of altitudes).",
+                f"Given triangle ABC with coordinates A={{A}}, B={{B}}, and C={{C}}, find the coordinates of its orthocenter. {GEOMETRY_FORMAT_INSTRUCTIONS}",
+                f"For triangle with vertices A={{A}}, B={{B}}, and C={{C}}, determine the orthocenter (intersection of altitudes). {GEOMETRY_FORMAT_INSTRUCTIONS}",
             ],
             "incircle_radius": [
-                "Consider triangle ABC with coordinates A={A}, B={B}, and C={C}. Compute the radius of its incircle.",
-                "Find the incircle radius of triangle ABC whose vertices are A={A}, B={B}, and C={C}.",
+                f"Consider triangle ABC with coordinates A={{A}}, B={{B}}, and C={{C}}. Compute the radius of its incircle. {GEOMETRY_FORMAT_INSTRUCTIONS}",
+                f"Find the incircle radius of triangle ABC whose vertices are A={{A}}, B={{B}}, and C={{C}}. {GEOMETRY_FORMAT_INSTRUCTIONS}",
             ],
             "angle_measure": [
-                "In triangle ABC with coordinates A={A}, B={B}, and C={C}, find the measure (in degrees) of angle ABC.",
-                "Given a triangle with vertices A={A}, B={B}, C={C}, determine the angle at B in degrees.",
+                f"In triangle ABC with coordinates A={{A}}, B={{B}}, and C={{C}}, find the measure (in degrees) of angle ABC. {GEOMETRY_FORMAT_INSTRUCTIONS}",
+                f"Given a triangle with vertices A={{A}}, B={{B}}, and C={{C}}, determine the angle at B in degrees. {GEOMETRY_FORMAT_INSTRUCTIONS}",
             ],
         }
         super().__init__(config=config, seed=config.seed, size=config.size)
@@ -77,6 +91,8 @@ class AdvancedGeometryDataset(ProceduralDataset):
         else:
             raise ValueError(f"Unknown task_type: {task_type}")
 
+        metadata["task_type"] = task_type
+
         return {
             "question": question,
             "answer": answer,
@@ -127,13 +143,14 @@ class AdvancedGeometryDataset(ProceduralDataset):
         y_ortho_approx = float(ortho.y.evalf())
 
         question_template = rng.choice(self._prompt_templates["orthocenter"])
-        question = question_template.format(A=(A.x, A.y), B=(B.x, B.y), C=(C.x, C.y))
+        question = question_template.format(A=(A.x, A.y), B=(B.x, B.y), C=(C.x, C.y), a="a", b="b")
         answer_str = f"({x_ortho_approx:.3f}, {y_ortho_approx:.3f})"
 
         metadata = {
             "A": (A.x, A.y),
             "B": (B.x, B.y),
             "C": (C.x, C.y),
+            "ortho": (ortho.x, ortho.y),
             "orthocenter_exact": (str(ortho.x), str(ortho.y)),
             "orthocenter_approx": (x_ortho_approx, y_ortho_approx),
         }
@@ -200,7 +217,7 @@ class AdvancedGeometryDataset(ProceduralDataset):
             angle_deg = float(angle_rad.evalf() * 180 / sympy.pi)
 
         question_template = rng.choice(self._prompt_templates["angle_measure"])
-        question = question_template.format(A=(A.x, A.y), B=(B.x, B.y), C=(C.x, C.y))
+        question = question_template.format(A=(A.x, A.y), B=(B.x, B.y), C=(C.x, C.y), a="a", b="b")
 
         answer_str = f"{angle_deg:.2f}°"
         metadata = {
@@ -211,6 +228,55 @@ class AdvancedGeometryDataset(ProceduralDataset):
         }
         return question, answer_str, metadata
 
+    def score_answer(self, answer: str | None, entry: Dict[str, Any]) -> float:
+        reward = 0.0
+        expected_answer = entry["answer"]
+        metadata = entry["metadata"]
+        task_type = metadata["task_type"]
+
+        if answer is not None:
+            try:
+                if metadata["task_type"] == "angle_measure":
+                    answer = answer.replace("\u00b0", "")
+                    expected_answer = expected_answer.replace("\u00b0", "")
+                    if np.round(float(answer), 2) == np.round(float(expected_answer), 2):
+                        reward = 1.0
+                    elif len(answer.strip()) > 0:
+                        reward = 0.05
+                    else:
+                        reward = 0.01
+                elif metadata["task_type"] == "orthocenter":
+                    x_coord = answer.split(",")[0].replace("(", "").strip()
+                    y_coord = answer.split(",")[1].replace(")", "").strip()
+
+                    expected_x = metadata["ortho"][0]
+                    expected_y = metadata["ortho"][1]
+
+                    if x_coord == expected_x and y_coord == expected_y:
+                        reward = 1.0
+                    elif (np.round(float(x_coord), 2) == np.round(float(expected_x), 2)) and (
+                        np.round(float(y_coord), 2) == np.round(float(expected_y), 2)
+                    ):
+                        reward = 1.0
+                    elif len(x_coord.strip()) > 0 and len(y_coord.strip()) > 0:
+                        reward = 0.05
+                    else:
+                        reward = 0.01
+                elif metadata["task_type"] == "incircle_radius":
+                    if answer == expected_answer:
+                        reward = 1.0
+                    elif np.round(float(answer), 2) == np.round(float(metadata["incircle_radius_exact"]), 2):
+                        reward = 0.5
+                    elif len(answer.strip()) > 0:
+                        reward = 0.05
+                    else:
+                        reward = 0.01
+                else:
+                    raise ValueError(f"Unknown task_type: {task_type}")
+            except:
+                reward = 0.01
+        return reward
+
 
 # Register the dataset
 register_dataset("advanced_geometry", AdvancedGeometryDataset, AdvancedGeometryConfig)

reasoning_gym/geometry/simple_geometry.py

@@ -46,15 +46,21 @@ class SimpleGeometryDataset(ProceduralDataset):
             (
                 "Given a convex polygon with {n_sides} sides, its first {n_minus_1} interior angles "
                 "are: {angle_list}. What is the measure of the remaining interior angle (in degrees)?"
+                "Return only the angle as your answer."
+                "Do not give the units in your answer."
             ),
             (
                 "A convex polygon has {n_sides} sides. The measures of "
                 "the first {n_minus_1} interior angles are: {angle_list}. "
                 "Find the measure of the last interior angle."
+                "Return only the angle as your answer."
+                "Do not give the units in your answer."
             ),
             (
                 "Consider a convex {n_sides}-gon whose first {n_minus_1} interior angles "
                 "are: {angle_list}. Determine the measure of the remaining angle."
+                "Return only the angle as your answer."
+                "Do not give the units in your answer."
             ),
         ]
         super().__init__(config=config, seed=config.seed, size=config.size)

reasoning_gym/logic/zebra_puzzles.py

@@ -44,6 +44,7 @@ class ZebraDataset(ProceduralDataset):
         q = instance["questions"][0]["question"]
         answer = instance["questions"][0]["answer"]
         question = str(puzzle) + "\n" + q
+        question = question + "? Provide only the name of the person as your final answer."
 
         return {
             "question": question,