Repository: open-thought/reasoning-gym

Issue:
Add score_answer() for PolynomialEquationsDataset
1. Override the `score_answer()` method in the `PolynomialEquationsDataset` class, parse the answer and calculate a reward depending on the absolute distance of the submitted answer and the ground truth correct solutions.
2. Add a hint to the question about the expected precision of the answer (e.g. number of decimal digits)
3. for the value of "answer" return all correct answers in a single string in the form which would score 1.0 when passed to score_answer()

See:
- [current output examples](https://github.com/open-thought/reasoning-gym/blob/main/GALLERY.md#polynomial_equations) of dataset in gallery
- ProceduralDataset base implementation of [score_answer()](https://github.com/open-thought/reasoning-gym/blob/8b0f634f4c75c0560b00fbbc0b0ec59e762a9bf2/reasoning_gym/dataset.py#L54-L66)

Discussion:
@rishabhranawat: @andreaskoepf is anyone working on this? If not, I'd like to try!

Clarification question: What's an appropriate score when the number of solutions are mismatched? 
   * Should return reward as 0.0
   * Or Treat the predicted answer as 0.0 and compute the reward

Computing the score plan:
1. Matching function: 
We may need to do some closest matching logic. In the case that the right_answers, predicted_answers match in the number of elements, just sorting should be enough. Else, we may need to some closest matching answer and then compute the score. Let's assume we have `matched_answers = f(right_answers, predicted_answers)`

2. Scoring function:
We use a exp decay for the score computation and normalize by the number of right answers to keep the reward between 0.0 and 1.0.
```
total_score = 0.0
for actual, predicted in zip(right_answers, matched_answers):
  distance  = abs(actual - predicted)
  total_score += exp^{-distance}

return total_score / len(right_answers)
```

WDYT?
@andreaskoepf: I would try: 
For each submitted answer search the closest `unmatched_right_answer`, add to total_score based on abs-distance as you suggested and remove entry from `unmatched_right_answer` (each can only be matched once). Divide by len(right_answers) in the end (as shown in your pseudo-code).@andreaskoepf: @rishabhranawat I assigned the task to you. Please let me know in case you don't have time to do it.@andreaskoepf: @rishabhranawat please let me know if you have time over the weekend to implement the `score_answer()` for PolynomialEquationsDataset .. otherwise please unassigned yourself.@rishabhranawat: Hi @andreaskoepf, plan to implement it this weekend, sorry for the delay in response.

Solution patch:
diff --git a/reasoning_gym/algebra/polynomial_equations.py b/reasoning_gym/algebra/polynomial_equations.py
index ed7e857..a182295 100644
--- a/reasoning_gym/algebra/polynomial_equations.py
+++ b/reasoning_gym/algebra/polynomial_equations.py
@@ -1,7 +1,8 @@
+import math
 import random
 import string
 from dataclasses import dataclass
-from typing import Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 from sympy import Eq, Symbol, expand, solve
 
@@ -26,6 +27,9 @@ class PolynomialEquationsConfig:
     )  # Allowed operators between terms, Avoid adding '*' or '/' because they will affect the degree
     seed: Optional[int] = None
     size: int = 500
+    # reward function hyperparameters
+    penalty_missing_factor = 0.1
+    penalty_extra_factor = 0.05
 
     def validate(self) -> None:
         """Validate configuration parameters."""
@@ -146,5 +150,101 @@ class PolynomialEquationsDataset(ProceduralDataset):
 
         return polynomial_expr
 
+    def _parse_score_to_list(self, answer: Optional[str]) -> List[float]:
+        """Parses a comma-separated string of scores into a sorted list of floats.
+
+        This method takes a string containing comma-separated numeric values,
+        attempts to convert each value to a float, and returns a sorted list of these floats.
+        Any values that cannot be converted to a float are ignored.
+        Handles empty strings gracefully.
+
+        Args:
+            answer: An optional string containing comma-separated numeric values.
+            Can be None or an empty string.
+        Returns:
+            A sorted list of floats parsed from the input string.
+            Returns an empty list if the input is None, empty, or contains no valid numeric values.
+        """
+
+        if answer is None or len(answer) == 0:  # Handle None or empty input
+            return []
+
+        output_float_vals = []
+        for output_val in answer.split(","):
+            try:
+                # Convert to float, strip whitespace
+                output_float_vals.append(float(output_val.strip()))
+            except ValueError:
+                # Ignore values that cannot be converted to float
+                continue
+
+        return sorted(output_float_vals)  # Return the sorted list of floats
+
+    def score_answer(self, answer: Optional[str], entry: Dict[str, any]) -> float:
+        """
+        Score an answer based on its numerical distance to oracle solutions using exponential decay.
+        This function compares a predicted answer (or list of answers) to a set of oracle solutions
+        (also a list of numbers). It calculates a reward based on how close the predicted solutions
+        are to the oracle solutions, using an exponential decay function.  It also applies penalties
+        for missing or extra predicted solutions. The implementation is a greedy algorithm where we
+        find the closest matching oracle solution for a given predicted solution and only allow an
+        oracle solution to match once.
+
+        Args:
+            answer: The predicted answer (or a string that can be parsed into a list of numbers).
+                    May be None.
+            entry: A dictionary containing the oracle solution(s) under the key "answer"
+                (which can be a string that can be parsed into a list of numbers).
+
+        Returns:
+            A float representing the final score. The score is non-negative.
+        """
+        oracle_solutions = self._parse_score_to_list(entry["answer"])  # Parse oracle solutions
+        predicted_solutions = self._parse_score_to_list(answer)  # Parse predicted solutions
+
+        total_reward = 0.0
+        matched_solutions = 0
+        extra_solutions = 0
+        missing_solutions = 0
+
+        for predicted_solution in predicted_solutions:
+
+            # find the closest matching solution from the oracle solutions.
+            # this is a greedy approach to computing the score
+            matched_distance = float("inf")
+            matched_distance_index = None
+            for oracle_solution_index, oracle_solution in enumerate(oracle_solutions):
+                if matched_distance > abs(predicted_solution - oracle_solution):
+                    matched_distance = abs(predicted_solution - oracle_solution)
+                    matched_distance_index = oracle_solution_index
+
+            if matched_distance_index is not None:
+                matched_solutions += 1
+                # Remove matched oracle solution
+                oracle_solutions.pop(matched_distance_index)
+                # Exponential decay reward
+                total_reward += math.exp(-matched_distance)
+            else:
+                # Extra predicted solution
+                extra_solutions += 1
+
+        # Count remaining oracle solutions as missing
+        for oracle_solution in oracle_solutions:
+            missing_solutions += 1
+
+        # Calculate penalty for either missing or extra solutions
+        penalty = missing_solutions * self.config.penalty_missing_factor
+        penalty += extra_solutions * self.config.penalty_extra_factor
+
+        if matched_solutions > 0:
+            # normalize the rewards that we found matching solutions for
+            # so that the value is bounded between 0 and 1
+            total_reward = total_reward / matched_solutions
+
+        # Final reward capped at 0
+        final_reward = max(0, total_reward - penalty)
+
+        return final_reward
+
 
 register_dataset("polynomial_equations", PolynomialEquationsDataset, PolynomialEquationsConfig)
diff --git a/reasoning_gym/arithmetic/prime_factorization.py b/reasoning_gym/arithmetic/prime_factorization.py
index d8c9d2a..55ec14a 100644
--- a/reasoning_gym/arithmetic/prime_factorization.py
+++ b/reasoning_gym/arithmetic/prime_factorization.py
@@ -1,8 +1,9 @@
 """Prime factorization task generator"""
 
+import math
 from dataclasses import dataclass
 from random import Random
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 from ..factory import ProceduralDataset, register_dataset
 
@@ -43,6 +44,25 @@ class PrimeFactorizationDataset(ProceduralDataset):
                 break
         return factors
 
+    def _normalize_answer(self, answer: str) -> List[int]:
+        """Parse and sort factors from a string"""
+        return sorted([int(factor.strip()) for factor in answer.split("×")])
+
+    def score_answer(self, answer: Optional[str], entry: Dict[str, any]) -> float:
+        oracle_answer = entry["answer"]
+        reward = 0.0
+        if answer is not None:
+            oracle_answer_parsed = self._normalize_answer(oracle_answer)
+            answer_parsed = self._normalize_answer(answer)
+            if oracle_answer_parsed == answer_parsed:
+                reward = 1.0
+            elif math.prod(oracle_answer_parsed) == math.prod(answer_parsed):
+                reward = 0.5
+            else:
+                reward = 0.01
+
+        return reward
+
     def __getitem__(self, idx: int) -> dict:
         """Generate a single prime factorization task"""
         rng = Random(self.seed + idx)


Test patch:
diff --git a/tests/test_polynomial_equations.py b/tests/test_polynomial_equations.py
index 6e1bb0c..e7caf65 100644
--- a/tests/test_polynomial_equations.py
+++ b/tests/test_polynomial_equations.py
@@ -1,4 +1,5 @@
 import pytest
+from pytest import approx
 from sympy import Symbol, sympify
 
 from reasoning_gym import create_dataset
@@ -115,3 +116,25 @@ def test_polynomial_solutions_evaluation():
                 f"Solution {solution} does not satisfy the polynomial {poly_str}. "
                 f"Evaluated value: {evaluated_value}"
             )
+
+
+@pytest.mark.parametrize(
+    "oracle_answer, predicted_answer, expected_reward",
+    [
+        ("4,-4.12", "4,-4.12", 1.0),  # Exact match
+        ("4,-4.12", "4.0001,-4.120001", approx(0.9999, rel=1e-3)),  # Very close match
+        ("4,-4.12", "4.1,-4.2", approx(0.9139, rel=1e-3)),
+        ("4,8", "4", approx(0.9, rel=1e-3)),  # Missing an oracle solution -> missing solution penalty applies
+        ("4", "4,8", approx(0.95, rel=1e-3)),  # extra solution -> extra solution penalty
+        ("-1,-2", "1,4", approx(0.06890, rel=1e-3)),  # -1 matched w/ 1 and -2 matched w/ 4
+        ("", "1", approx(0, rel=1e-4)),  # oracle no solution, predicted extra solution
+        ("1", "", approx(0, rel=1e-4)),  # oracle has a solution, predicted no solution
+    ],
+)
+def test_polynomial_solutions_score_answer(oracle_answer, predicted_answer, expected_reward):
+    # You might want to parameterize cfg as well
+    cfg = PolynomialEquationsConfig(seed=999, size=3)
+    ds = PolynomialEquationsDataset(cfg)
+
+    actual_reward = ds.score_answer(predicted_answer, {"answer": oracle_answer})
+    assert actual_reward == pytest.approx(expected_reward, rel=1e-3)  # Fuzzy comparison for floats
diff --git a/tests/test_prime_factorization.py b/tests/test_prime_factorization.py
index b3b8695..70c463b 100644
--- a/tests/test_prime_factorization.py
+++ b/tests/test_prime_factorization.py
@@ -85,6 +85,37 @@ def test_prime_factorization_known_values():
     assert item["answer"] == "2 × 2 × 3"
 
 
+def test_prime_factorization_score_answer():
+    """Test scoring of answers"""
+    config = PrimeFactorizationConfig(min_value=12, max_value=12, size=1, seed=42)  # Force specific number
+    dataset = PrimeFactorizationDataset(config)
+    item = dataset[0]
+
+    # Perfectly ordered answer
+    answer = "2 × 2 × 3"
+    assert dataset.score_answer(answer, item) == 1.0
+
+    # No white spaces answer (still correct)
+    answer = "2×2×3"
+    assert dataset.score_answer(answer, item) == 1.0
+
+    # Shuffled factors (still correct)
+    answer = "2 × 3 × 2"
+    assert dataset.score_answer(answer, item) == 1.0
+
+    # Partially correct answer (not all numbers are fully factorized)
+    answer = "2 × 6"
+    assert dataset.score_answer(answer, item) == 0.5
+
+    # Incorrect answer
+    answer = "2 × 5"
+    assert dataset.score_answer(answer, item) == 0.01
+
+    # Answer is none
+    answer = None
+    assert dataset.score_answer(answer, item) == 0.0
+
+
 def is_prime(n: int) -> bool:
     """Helper function to check if a number is prime"""
     if n < 2: