Merge branch 'main' into rich/graphcolor

reasoning_gym/algebra/intermediate_integration.py

@@ -76,6 +76,11 @@ class IntermediateIntegrationDataset(ProceduralDataset):
             "Calculate the antiderivative: ∫ {integrand} dx",
             "Evaluate the indefinite integral: ∫ {integrand} dx",
         ]
+        self.added_instruction = """
+In addition, when doing calculation, use the following instructions together with your mathematical ingenuity to solve the integral problems
+## 1. Use ** instead ^ to represent powers. For example 7*X**2 instead of 7*X^2.
+## 2. Always use * when doing all sorts of multiplcation in your reasoning steps. For example Use [-3*X**3*sin(X) - 9*X**2*cos(X) + 18*X*sin(X) + 18*cos(X) + C] instead of [-3x3sin(x) - 9x2cos(x) + 18xsin(x) + 18cos(x) + C].
+"""
 
     def _get_outer_constant(self, rng: random.Random) -> int:
         """Helper to generate signed outer constant from config"""
@@ -222,9 +227,10 @@ class IntermediateIntegrationDataset(ProceduralDataset):
 
         answer = sympy.integrate(integrand, x)
         answer_str = str(answer) + " + C"
+        question = rng.choice(self.prompt_template).format(integrand=integrand) + self.added_instruction
 
         return {
-            "question": rng.choice(self.prompt_template).format(integrand=integrand),
+            "question": question,
             "answer": answer_str,
             "metadata": {
                 "integrand": str(integrand),

reasoning_gym/algebra/polynomial_equations.py

@@ -62,6 +62,14 @@ class PolynomialEquationsDataset(ProceduralDataset):
             "Determine the real value(s) of {variable} that satisfies: {polynomial_expanded} = 0",
             "Solve the polynomial equation for real {variable}:\n{polynomial_expanded} = 0",
         ]
+        self.added_instruction = """
+In solving the equations, please abide by the following instruction:
+## 1. All answers should be comma-separated. For example "-0.3773, 0.4005" etc.
+## 2. In cases where your answer is b = 2 + sqrt(4560) / 172 and b = 2 - sqrt(4560) / 172. Since b can be 2 numbers, resolve your answer like this instead, "-0.3773, 0.4005".
+## 3. If there are no real values of i that satisfy the equation, report your answer as empty string, "".
+## 4. If there are 2 answers, resolve the answers as comma-separated floats of 2 numbers, if 3 answers, make it comma-separated floats of 3 numbers.
+## 5. Resolve all numbers as floats in the string of comma-separated numbers. Round the floats higher than 4 decimal place(d.p) down to 4 d.p.
+"""
         super().__init__(config=config, seed=config.seed, size=config.size)
 
     def __getitem__(self, idx: int) -> dict:
@@ -89,19 +97,20 @@ class PolynomialEquationsDataset(ProceduralDataset):
             for sol in solutions:
                 if sol.is_real:
                     # Evaluate symbolic solution to a floating approximation
-                    real_solutions.append(float(sol.evalf()))
+                    real_solutions.append(round(float(sol.evalf()), 4))
 
             if len(real_solutions) > 0:
                 real_solutions.sort()
                 break
 
         answer_str = ", ".join(str(x) for x in real_solutions)
+        question = (
+            rng.choice(self._prompt_templates).format(variable=variable, polynomial_expanded=polynomial_expanded)
+            + self.added_instruction
+        )
 
         return {
-            "question": rng.choice(self._prompt_templates).format(
-                variable=variable,
-                polynomial_expanded=polynomial_expanded,
-            ),
+            "question": question,
             "answer": answer_str,
             "metadata": {
                 "polynomial_expr": str(polynomial_expanded),

reasoning_gym/algebra/polynomial_multiplication.py

@@ -61,6 +61,11 @@ class PolynomialMultiplicationDataset(ProceduralDataset):
             "Simplify this expression: {polynomial_expr}",
             "Calculate the following: {polynomial_expr}",
         ]
+        self.added_instruction = """
+In addition, When doing calculation, Use the following instructions together with your mathematical ingenuity to solve the integral problems
+## 1. Use ** instead ^ to represent powers. For example 7*X**2 instead of 7*X^2.
+## 2. Always use * when doing all sorts of multiplcation in your reasoning steps and even in reporting answers.
+"""
         super().__init__(config=config, seed=config.seed, size=config.size)
 
     def __getitem__(self, idx: int) -> dict:
@@ -79,11 +84,10 @@ class PolynomialMultiplicationDataset(ProceduralDataset):
 
         polynomial_expr = sp.prod(polynomials)
         product = sp.expand(polynomial_expr)
+        question = rng.choice(self._prompt_templates).format(polynomial_expr=polynomial_expr) + self.added_instruction
 
         return {
-            "question": rng.choice(self._prompt_templates).format(
-                polynomial_expr=polynomial_expr,
-            ),
+            "question": question,
             "answer": product,
             "metadata": {
                 "polynomial_expr": str(polynomial_expr),

reasoning_gym/algebra/simple_integration.py

@@ -41,6 +41,11 @@ class SimpleIntegrationDataset(ProceduralDataset):
             "Calculate the antiderivative: ∫ {integrand} dx",
             "Evaluate the indefinite integral: ∫ {integrand} dx",
         ]
+        self.added_instruction = """
+In addition, When doing calculation, Use the following instructions together with your mathematical ingenuity to solve the integral problems
+## 1. Use ** instead ^ to represent powers. For example 7*X**2 instead of 7*X^2.
+## 2. Always use * when doing all sorts of multiplcation in your reasoning steps. For example Use [-3*X**3*sin(X) - 9*X**2*cos(X) + 18*X*sin(X) + 18*cos(X) + C] instead of [-3x3sin(x) - 9x2cos(x) + 18xsin(x) + 18cos(x) + C].
+"""
         super().__init__(config=config, seed=config.seed, size=config.size)
 
     def _generate_coefficient(self, rng: random.Random) -> Fraction:
@@ -69,9 +74,10 @@ class SimpleIntegrationDataset(ProceduralDataset):
         rng = random.Random(self.seed + idx)
         symbol, polynomial = self._generate_polynomial(rng)
         derivative = sympy.diff(polynomial, symbol)
+        question = rng.choice(self._prompt_templates).format(integrand=derivative) + self.added_instruction
 
         return {
-            "question": rng.choice(self._prompt_templates).format(integrand=derivative),
+            "question": question,
             "answer": str(polynomial) + " + C",
             "metadata": {
                 "integrand": str(derivative),

reasoning_gym/algorithmic/__init__.py

@@ -26,7 +26,6 @@ from .rotate_matrix import RotateMatrixConfig, RotateMatrixDataset
 from .sentence_reordering import SentenceReorderingConfig, SentenceReorderingDataset
 from .spell_backward import SpellBackwardConfig, SpellBackwardDataset
 from .spiral_matrix import SpiralMatrixConfig, SpiralMatrixDataset
-from .string_insertion import StringInsertionConfig, StringInsertionDataset
 from .string_manipulation import StringManipulationConfig, StringManipulationDataset
 from .word_ladder import WordLadderConfig, WordLadderDataset
 from .word_sequence_reversal import WordSequenceReversalConfig, WordSequenceReversalDataset

reasoning_gym/algorithmic/number_sorting.py

@@ -34,6 +34,11 @@ class NumberSortingDataset(ProceduralDataset):
 
     def __init__(self, config: NumberSortingConfig):
         super().__init__(config=config, seed=config.seed, size=config.size)
+        self.added_instruction = """
+Please follow the instruction below:
+## 1. Let all your answers be a list of numbers. Instead of reporting your answer as -69, -13, 1, 7, 11, 43, 59, 61, use ['-69', '-13', '1', '7', '11', '43', '59', '61'] instead
+## 2. Convert all numbers in the square brackets as strings. For example, ['-69', '-13', '1', '7', '11', '43', '59', '61']
+"""
 
     def _format_number(self, num: float, decimals: int) -> str:
         """Format number with specified decimal places"""
@@ -78,9 +83,10 @@ class NumberSortingDataset(ProceduralDataset):
         is_ascending = rng.choice([True, False])
         direction = "ascending" if is_ascending else "descending"
         answer = asc_answer if is_ascending else desc_answer
+        question = f"Sort these numbers in {direction} order: {', '.join(number_strs)}" + self.added_instruction
 
         return {
-            "question": f"Sort these numbers in {direction} order: {', '.join(number_strs)}",
+            "question": question,
             "answer": str(answer),
             "metadata": {"original_numbers": number_strs, "direction": direction, "sorted_numbers": answer},
         }

reasoning_gym/arc/arc_1d.py

@@ -58,27 +58,27 @@ class Arc1DDataset(ProceduralDataset):
                 - metadata: dict with generation parameters
         """
         # Create deterministic RNG from base seed and idx
-        item_rng = Random(self.seed + idx)
+        rng = Random(self.seed + idx)
 
         # Select random task
-        task_name = item_rng.choice(self.task_names)
+        task_name = rng.choice(self.task_names)
         task_func, task_kwargs = self.ARC_1D_TASKS[task_name]
 
         # Generate training examples
         train_examples = []
-        size = item_rng.randint(self.config.min_size, self.config.max_size)
+        size = rng.randint(self.config.min_size, self.config.max_size)
 
         for _ in range(self.config.num_train):
             example = None
             while example is None:
-                example = task_func(item_rng, size, **task_kwargs)
+                example = task_func(rng, size, **task_kwargs)
 
             train_examples.append(example)
 
         # Generate test example
         test_example = None
         while test_example is None:
-            test_example = task_func(item_rng, size, **task_kwargs)
+            test_example = task_func(rng, size, **task_kwargs)
 
         # Format question
         question = "Find the common rule that maps an input grid to an output grid, given the examples below.\n\n"

reasoning_gym/arithmetic/__init__.py

@@ -4,7 +4,7 @@ Arithmetic tasks for training reasoning capabilities:
 
 from .basic_arithmetic import BasicArithmeticDataset, BasicArithmeticDatasetConfig
 from .calendar_arithmetic import CalendarArithmeticConfig, CalendarArithmeticDataset
-from .chain_sum import ChainSum, ChainSumConfig
+from .chain_sum import ChainSumConfig, ChainSumDataset
 from .count_bits import CountBitsConfig, CountBitsDataset
 from .dice import DiceConfig, DiceDataset
 from .fraction_simplification import FractionSimplificationConfig, FractionSimplificationDataset
@@ -14,12 +14,13 @@ from .lcm import LCMConfig, LCMDataset
 from .leg_counting import LegCountingConfig, LegCountingDataset
 from .power_function import PowerFunctionConfig, PowerFunctionDataset
 from .prime_factorization import PrimeFactorizationConfig, PrimeFactorizationDataset
+from .products import ProductsConfig, ProductsDataset
 from .time_intervals import TimeIntervalsConfig, TimeIntervalsDataset
 
 __all__ = [
     "BasicArithmeticDataset",
     "BasicArithmeticDatasetConfig",
-    "ChainSum",
+    "ChainSumDataset",
     "ChainSumConfig",
     "CalendarArithmeticConfig",
     "CalendarArithmeticDataset",
@@ -31,8 +32,12 @@ __all__ = [
     "LCMDataset",
     "LegCountingConfig",
     "LegCountingDataset",
+    "PowerFunctionConfig",
+    "PowerFunctionDataset",
     "PrimeFactorizationConfig",
     "PrimeFactorizationDataset",
+    "ProductsDataset",
+    "ProductsConfig",
     "GSMSymbolicDatasetConfig",
     "GSMSymbolicDataset",
     "TimeIntervalsConfig",

reasoning_gym/arithmetic/basic_arithmetic.py

@@ -78,17 +78,17 @@ class BasicArithmeticDataset(ProceduralDataset):
                 - metadata: dict with generation parameters
         """
         # Create deterministic RNG from base seed and idx
-        item_rng = Random(self.seed + idx)
+        rng = Random(self.seed + idx)
 
-        num_terms = item_rng.randint(self.config.min_terms, self.config.max_terms)
-        num_digits = item_rng.randint(self.config.min_digits, self.config.max_digits)
+        num_terms = rng.randint(self.config.min_terms, self.config.max_terms)
+        num_digits = rng.randint(self.config.min_digits, self.config.max_digits)
 
         if self.config.allow_parentheses:
-            expression, result = self._generate_complex_task(item_rng, num_terms, num_digits)
+            expression, result = self._generate_complex_task(rng, num_terms, num_digits)
         else:
-            expression, result = self._generate_simple_task(item_rng, num_terms, num_digits)
+            expression, result = self._generate_simple_task(rng, num_terms, num_digits)
 
-        question = self._format_question(item_rng, expression)
+        question = self._format_question(rng, expression)
 
         return {
             "question": question,

reasoning_gym/arithmetic/calendar_arithmetic.py

@@ -122,9 +122,9 @@ class CalendarArithmeticDataset(ProceduralDataset):
         self.tasks = [self.task_handlers[task] for task in self.config.tasks]
 
     def __getitem__(self, idx: int) -> dict:
-        item_rng = random.Random(self.seed + idx)
-        task = item_rng.choice(self.tasks)
-        question, answer, metadata = task(item_rng)
+        rng = random.Random(self.seed + idx)
+        task = rng.choice(self.tasks)
+        question, answer, metadata = task(rng)
         return {
             "question": question,
             "answer": str(answer),

reasoning_gym/arithmetic/chain_sum.py

@@ -32,7 +32,7 @@ class ChainSumConfig:
             assert 10 ** (self.min_digits - 1) >= 1, "min_digits would result in invalid number range"
 
 
-class ChainSum(ProceduralDataset):
+class ChainSumDataset(ProceduralDataset):
     """Generates simple arithmetic tasks using only + and - operators"""
 
     def __init__(self, config: ChainSumConfig):
@@ -51,16 +51,16 @@ class ChainSum(ProceduralDataset):
                 - metadata: dict with generation parameters
         """
         # Create deterministic RNG from base seed and idx
-        item_rng = random.Random(self.seed + idx)
+        rng = random.Random(self.seed + idx)
 
-        num_terms = item_rng.randint(self.config.min_terms, self.config.max_terms)
-        num_digits = item_rng.randint(self.config.min_digits, self.config.max_digits)
+        num_terms = rng.randint(self.config.min_terms, self.config.max_terms)
+        num_digits = rng.randint(self.config.min_digits, self.config.max_digits)
 
         # Calculate value ranges based on number of digits
         min_value = 0 if num_digits == 1 else 10 ** (num_digits - 1)  # Special case for 1 digit
         max_value = (10**num_digits) - 1  # e.g., 999 for 3 digits
 
-        expression, result = self._generate_task(item_rng, num_terms, min_value, max_value)
+        expression, result = self._generate_task(rng, num_terms, min_value, max_value)
 
         return {
             "question": f"{expression} =",
@@ -143,4 +143,4 @@ class ChainSumCurriculum(BaseCurriculum):
 
 
 # Register the dataset
-register_dataset("chain_sum", ChainSum, ChainSumConfig)
+register_dataset("chain_sum", ChainSumDataset, ChainSumConfig)

reasoning_gym/arithmetic/products.py

@@ -0,0 +1,130 @@
+import random
+from dataclasses import dataclass
+from typing import Optional
+
+from ..coaching import AttributeType, BaseCurriculum, RangeAttributeDefinition
+from ..factory import ProceduralDataset, register_dataset
+
+
+@dataclass
+class ProductsConfig:
+    """Configuration for products task generation"""
+
+    min_terms: int = 2
+    max_terms: int = 2
+    min_digits: int = 1
+    max_digits: int = 5
+    seed: Optional[int] = None
+    size: int = 500
+
+    def validate(self) -> None:
+        """Validate configuration parameters"""
+        assert self.size > 0, "size must be positive"
+        assert self.min_terms > 0, "min_terms must be positive"
+        assert self.max_terms >= self.min_terms, "max_terms must be >= min_terms"
+        assert self.min_digits > 0, "min_digits must be positive"
+        assert self.max_digits >= self.min_digits, "max_digits must be >= min_digits"
+
+
+class ProductsDataset(ProceduralDataset):
+    """Generates multiplication tasks with configurable number of terms"""
+
+    def __init__(self, config: ProductsConfig):
+        super().__init__(config=config, seed=config.seed, size=config.size)
+
+    def __getitem__(self, idx: int) -> dict:
+        """Generate a single multiplication task
+
+        Args:
+            idx: Index of the item to generate
+
+        Returns:
+            dict with keys:
+                - question: str, the formatted multiplication expression
+                - answer: str, the ground truth result
+                - metadata: dict with generation parameters
+        """
+        # Create deterministic RNG from base seed and idx
+        rng = random.Random(self.seed + idx)
+
+        num_terms = rng.randint(self.config.min_terms, self.config.max_terms)
+        num_digits = rng.randint(self.config.min_digits, self.config.max_digits)
+
+        # Calculate value ranges based on number of digits
+        min_value = 0 if num_digits == 1 else 10 ** (num_digits - 1)  # Special case for 1 digit
+        max_value = (10**num_digits) - 1  # e.g., 999 for 3 digits
+
+        expression, result = self._generate_task(rng, num_terms, min_value, max_value)
+
+        return {
+            "question": f"{expression} =",
+            "answer": str(result),
+            "metadata": {
+                "difficulty": {
+                    "num_terms": num_terms,
+                    "num_digits": num_digits,
+                },
+                "expression": expression,
+            },
+        }
+
+    def _generate_task(self, rng: random.Random, num_terms: int, min_value: int, max_value: int) -> tuple[str, int]:
+        """Generate a multiplication task
+
+        Args:
+            rng: Random number generator
+            num_terms: Number of terms in the expression
+            min_value: Minimum value for generated numbers
+            max_value: Maximum value for generated numbers
+
+        Returns:
+            Tuple of (expression string, result integer)
+        """
+        # Generate random numbers within the specified range
+        constants = [rng.randint(min_value, max_value) for _ in range(num_terms)]
+
+        # Build expression and compute result
+        expression_parts = []
+        result = constants[0]
+
+        expression_parts.append(str(constants[0]))
+        for i in range(1, len(constants)):
+            expression_parts.append("*")
+            expression_parts.append(str(constants[i]))
+            result *= constants[i]
+
+        expression = " ".join(expression_parts)
+        return expression, result
+
+
+class ProductsCurriculum(BaseCurriculum):
+    def __init__(self):
+        super().__init__(ProductsCurriculum.__name__, ProductsConfig)
+
+        # Define attributes
+        self._define_attributes(
+            RangeAttributeDefinition(
+                name="num_terms",
+                levels=[2, 3, 4, 5],
+                default_level=0,  # Start with 2 terms
+                description="Maximum number of terms in the expression",
+                attr_type=AttributeType.APPEND,
+                min_value=2,  # Ensure at least 2 terms
+                lower_field_name="min_terms",
+                upper_field_name="max_terms",
+            ),
+            RangeAttributeDefinition(
+                name="num_digits",
+                levels=[1, 2, 3, 4],
+                default_level=0,  # Start with 1-digit numbers
+                description="Number of digits in each operand",
+                attr_type=AttributeType.APPEND,
+                min_value=1,  # Ensure numbers are at least 1 digit
+                lower_field_name="min_digits",
+                upper_field_name="max_digits",
+            ),
+        )
+
+
+# Register the dataset
+register_dataset("products", ProductsDataset, ProductsConfig)

reasoning_gym/arithmetic/time_intervals.py

@@ -82,14 +82,14 @@ class TimeIntervalsDataset(ProceduralDataset):
 
     def __getitem__(self, idx: int) -> dict:
         """Generate a single time interval calculation task"""
-        item_rng = random.Random(self.seed + idx)
+        rng = random.Random(self.seed + idx)
 
         # Randomly choose task type from config
-        task_type = item_rng.choice(self.config.task_types)
+        task_type = rng.choice(self.config.task_types)
 
-        start_time, end_time, format_str, expected_format = self._generate_times(item_rng, task_type)
+        start_time, end_time, format_str, expected_format = self._generate_times(rng, task_type)
 
-        template = item_rng.choice(self.TEMPLATES)
+        template = rng.choice(self.TEMPLATES)
         question = template.format(start=start_time, end=end_time, format=expected_format)
 
         # Calculate the actual difference

reasoning_gym/games/game_of_life.py

@@ -1,6 +1,7 @@
+import json
 from dataclasses import dataclass
 from random import Random
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, Optional
 
 import cellpylib as cpl
 
@@ -11,8 +12,8 @@ from ..factory import ProceduralDataset, register_dataset
 class GameOfLifeConfig:
     """Configuration for sudoku puzzle generation"""
 
-    grid_size_x: int = 20
-    grid_size_y: int = 20
+    grid_size_x: int = 10
+    grid_size_y: int = 10
     filled_cells: int = 100  # actually a max
     simulation_steps: int = 1
     seed: Optional[int] = None
@@ -31,7 +32,7 @@ class GameOfLifeDataset(ProceduralDataset):
 
     def __init__(self, config: GameOfLifeConfig):
         self._prompt_templates = [
-            "What will this Game of Life board look like after {simulation_steps} steps of simulation?\n\n{board}"
+            "What will this Game of Life board look like after {simulation_steps} steps of simulation? Reply as array of array representing rows in the grid from top to bottom in JSON format. (An empty 3x3 grid would look like this: [[0,0,0],[0,0,0],[0,0,0]])\n\n{board}."
         ]
 
         super().__init__(config=config, seed=config.seed, size=config.size)
@@ -59,11 +60,18 @@ class GameOfLifeDataset(ProceduralDataset):
 
         # Simulate the result to get the answer
         evolved = cpl.evolve2d(
-            board, timesteps=self.config.simulation_steps + 1, apply_rule=cpl.game_of_life_rule, memoize="recursive"
+            board,
+            timesteps=self.config.simulation_steps + 1,
+            apply_rule=cpl.game_of_life_rule,
+            memoize="recursive",
         )
 
-        board_str = str(board[0])
-        result_str = str(evolved[-1])
+        rows = [json.dumps(board[0, i].tolist(), separators=(",", ":")) for i in range(board.shape[1])]
+        board_str = "[" + ", \n ".join(rows) + "]"
+
+        final_step = evolved[-1]
+        final_step_list = final_step.tolist()
+        result_str = json.dumps(final_step_list, separators=(",", ":"))
 
         return {
             "question": rng.choice(self._prompt_templates).format(
@@ -93,10 +101,17 @@ class GameOfLifeDataset(ProceduralDataset):
 
         if answer == None:
             return 0.0
-        if answer.replace("\n", "") != entry["answer"].replace("\n", ""):
+
+        try:
+            ans_arr = json.loads(answer)
+            correct_arr = json.loads(entry["answer"])
+
+            if correct_arr != ans_arr:
+                return 0.01
+            else:
+                return 1.0  # Yay
+        except Exception as e:
             return 0.01
-        else:
-            return 1.0  # Yay
 
 
 register_dataset("game_of_life", GameOfLifeDataset, GameOfLifeConfig)