From ec3e414a8c1d0a4639310611ab1b5f0d8bcdce8b Mon Sep 17 00:00:00 2001
From: Adefioye <47661641+Adefioye@users.noreply.github.com>
Date: Thu, 13 Mar 2025 15:03:57 -0500
Subject: [PATCH] Cryptarithm curriculum (#346)

* Add curriculum for cryptarithm
* Add difficulty to metadata
---
 reasoning_gym/algorithmic/__init__.py    |  3 +-
 reasoning_gym/algorithmic/cryptarithm.py | 28 ++++++++++++++-
 tests/test_cryptarithm.py                | 44 +++++++++++++++++++++++-
 3 files changed, 72 insertions(+), 3 deletions(-)

diff --git a/reasoning_gym/algorithmic/__init__.py b/reasoning_gym/algorithmic/__init__.py
index 6620310a..241dd369 100644
--- a/reasoning_gym/algorithmic/__init__.py
+++ b/reasoning_gym/algorithmic/__init__.py
@@ -12,7 +12,7 @@ from .binary_alternation import BinaryAlternationConfig, BinaryAlternationCurric
 from .binary_matrix import BinaryMatrixConfig, BinaryMatrixCurriculum, BinaryMatrixDataset
 from .caesar_cipher import CaesarCipherConfig, CaesarCipherCurriculum, CaesarCipherDataset
 from .count_primes import CountPrimesConfig, CountPrimesCurriculum, CountPrimesDataset
-from .cryptarithm import CryptarithmConfig, CryptarithmDataset
+from .cryptarithm import CryptarithmConfig, CryptarithmCurriculum, CryptarithmDataset
 from .game_of_life import GameOfLifeConfig, GameOfLifeDataset
 from .game_of_life_halting import GameOfLifeHaltingConfig, GameOfLifeHaltingDataset
 from .graph_color import GraphColorConfig, GraphColorCurriculum, GraphColorDataset
@@ -61,6 +61,7 @@ __all__ = [
     "CaesarCipherCurriculum",
     "CryptarithmConfig",
     "CryptarithmDataset",
+    "CryptarithmCurriculum",
     "GameOfLifeConfig",
     "GameOfLifeDataset",
     "GameOfLifeHaltingConfig",
diff --git a/reasoning_gym/algorithmic/cryptarithm.py b/reasoning_gym/algorithmic/cryptarithm.py
index fe62b743..b3f7d767 100644
--- a/reasoning_gym/algorithmic/cryptarithm.py
+++ b/reasoning_gym/algorithmic/cryptarithm.py
@@ -15,6 +15,7 @@ from dataclasses import dataclass
 from random import Random
 from typing import Any, Optional
 
+from ..coaching import AttributeType, BaseCurriculum, RangeAttributeDefinition
 from ..factory import ProceduralDataset, register_dataset
 
 
@@ -185,6 +186,10 @@ class CryptarithmDataset(ProceduralDataset):
                 "result_letters": result_letters,
                 "digit_to_letter": digit_to_letter,
                 "letter_to_digit": letter_to_digit,
+                "difficulty": {
+                    "min_words": self.config.min_words,
+                    "max_words": self.config.max_words,
+                },
             },
         }
 
@@ -237,4 +242,25 @@ class CryptarithmDataset(ProceduralDataset):
         return (total_correct / total) * 0.7 + 0.3
 
 
-register_dataset("cryptarithm", CryptarithmDataset, CryptarithmConfig)
+class CryptarithmCurriculum(BaseCurriculum):
+    """Curriculum for Cryptarithm puzzles."""
+
+    def __init__(self):
+        super().__init__(CryptarithmCurriculum.__name__, CryptarithmConfig)
+
+        # Define the attributes
+        self._define_attributes(
+            RangeAttributeDefinition(
+                name="words",
+                levels=[2, 5, 10, 50],
+                default_level=1,
+                description="Number of words in the cryptarithm puzzle",
+                attr_type=AttributeType.APPEND,
+                min_value=1,
+                lower_field_name="min_words",
+                upper_field_name="max_words",
+            )
+        )
+
+
+register_dataset("cryptarithm", CryptarithmDataset, CryptarithmConfig, CryptarithmCurriculum)
diff --git a/tests/test_cryptarithm.py b/tests/test_cryptarithm.py
index 64b7b5e4..41bd247e 100644
--- a/tests/test_cryptarithm.py
+++ b/tests/test_cryptarithm.py
@@ -1,7 +1,7 @@
 import pytest
 
 from reasoning_gym import create_dataset
-from reasoning_gym.algorithmic.cryptarithm import CryptarithmConfig, CryptarithmDataset
+from reasoning_gym.algorithmic.cryptarithm import CryptarithmConfig, CryptarithmCurriculum, CryptarithmDataset
 
 
 def test_cryptarithm_generation():
@@ -167,3 +167,45 @@ def test_cryptarithm_score_answer():
     # The formula is (num_correct / total) * 0.7 + 0.3
     expected_score = (half / total) * 0.7 + 0.3
     assert abs(score - expected_score) < 1e-9, f"Partial correctness: expected {expected_score}, got {score}"
+
+
+def test_cryptarithm_curriculum():
+    """Test curriculum for cryptarithm dataset"""
+
+    curriculum = CryptarithmCurriculum()
+    base_value = {"size": 150, "seed": 1}
+
+    base_cfg: CryptarithmCurriculum = curriculum.generate_configuration(base_value)
+
+    assert base_cfg.seed == 1
+    assert base_cfg.size == 150
+    assert base_cfg.min_words == 2
+    assert base_cfg.max_words == 5
+
+    # Test and validate increase in level
+    curriculum.increment_attr_level("words")
+    increased_cfg: CryptarithmCurriculum = curriculum.generate_configuration(base_value)
+
+    assert increased_cfg.min_words == 2
+    assert increased_cfg.max_words == 10
+
+    # Test and validate decrease in level
+    curriculum.decrement_attr_level("words")
+    decreased_cfg: CryptarithmCurriculum = curriculum.generate_configuration(base_value)
+
+    assert decreased_cfg.min_words == 2
+    assert decreased_cfg.max_words == 5
+
+    # Test upper bound boundary conditions
+    for _ in range(10):
+        curriculum.increment_attr_level("words")
+    upper_bound_cfg: CryptarithmCurriculum = curriculum.generate_configuration(base_value)
+    assert upper_bound_cfg.min_words == 2
+    assert upper_bound_cfg.max_words == 50
+
+    # Test lower bound boundary conditions
+    for _ in range(10):
+        curriculum.decrement_attr_level("words")
+    lower_bound_cfg: CryptarithmCurriculum = curriculum.generate_configuration(base_value)
+    assert lower_bound_cfg.min_words == 2
+    assert lower_bound_cfg.max_words == 2