Repository: open-thought/reasoning-gym

Issue:
fix counting_primes dataset metadata
![Image](https://github.com/user-attachments/assets/224863e7-534e-4fc8-b470-f121c2637576)

Hints from issue comments:

Solution patch:
diff --git a/reasoning_gym/algorithmic/count_primes.py b/reasoning_gym/algorithmic/count_primes.py
index 0a553c7..e505b91 100644
--- a/reasoning_gym/algorithmic/count_primes.py
+++ b/reasoning_gym/algorithmic/count_primes.py
@@ -51,8 +51,8 @@ class CountPrimesDataset(ProceduralDataset):
         rng = Random(self.seed + idx)
         start = rng.randint(1, self.config.max_n)
         end = rng.randint(start, self.config.max_n)
-        primes = self.primes[start : end + 1]
-        answer = sum(primes)
+        primes = [i for i in range(start, end + 1) if self.primes[i]]
+        answer = len(primes)
         return {
             "question": QUESTION_TEMPLATE.format(start=start, end=end),
             "answer": str(answer),
diff --git a/reasoning_gym/cognition/color_cube_rotation.py b/reasoning_gym/cognition/color_cube_rotation.py
index fe7feee..59e3e3a 100644
--- a/reasoning_gym/cognition/color_cube_rotation.py
+++ b/reasoning_gym/cognition/color_cube_rotation.py
@@ -191,17 +191,12 @@ class ColorCubeRotationDataset(ProceduralDataset):
 
     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"]
+                answer_formatted = answer.strip().lower()
+                solved = answer_formatted == entry["answer"].strip().lower()
                 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:


Test patch:
diff --git a/tests/test_color_cube_rotation.py b/tests/test_color_cube_rotation.py
index a554afd..87ecd8c 100644
--- a/tests/test_color_cube_rotation.py
+++ b/tests/test_color_cube_rotation.py
@@ -49,6 +49,16 @@ def test_deterministic_generation():
         assert dataset1[i]["question"] == dataset2[i]["question"]
         assert dataset1[i]["answer"] == dataset2[i]["answer"]
 
+    for item in dataset1:
+        assert isinstance(item, dict)
+        assert "question" in item
+        assert "answer" in item
+        assert "metadata" in item
+
+        # Test the scoring
+        assert dataset1.score_answer(answer=item["answer"], entry=item) == 1.0
+        assert dataset1.score_answer(answer=None, entry=item) == 0.0
+
 
 def test_cube_rotations():
     # Test individual rotation operations
diff --git a/tests/test_count_primes.py b/tests/test_count_primes.py
index f131647..68c642f 100644
--- a/tests/test_count_primes.py
+++ b/tests/test_count_primes.py
@@ -86,3 +86,18 @@ def test_count_primes_answer():
     assert primes[8] == False
     assert primes[9] == False
     assert primes[10] == False
+
+
+def test_count_primes_list():
+    """Test that list of primes was correctly generated"""
+    config = CountPrimesConfig(max_n=100, size=100, seed=42)
+    dataset = CountPrimesDataset(config)
+
+    for item in dataset:
+        start = item["metadata"]["start"]
+        end = item["metadata"]["end"]
+        primes = item["metadata"]["primes"]
+        for p in primes:
+            assert p >= start
+            assert p <= end
+            assert dataset.primes[p] == True