feat(env): Count Primes Curriculum (#287)

* count primes curriculum

reasoning_gym/algorithmic/count_primes.py

@@ -9,6 +9,7 @@ from dataclasses import dataclass
 from random import Random
 from typing import Optional
 
+from ..coaching import AttributeType, BaseCurriculum, RangeAttributeDefinition
 from ..factory import ProceduralDataset, register_dataset
 
 QUESTION_TEMPLATE = """Count how many prime numbers there are between {start} and {end} (inclusive) ?"""
@@ -18,6 +19,7 @@ QUESTION_TEMPLATE = """Count how many prime numbers there are between {start} an
 class CountPrimesConfig:
     """Configuration for Count Primes dataset generation"""
 
+    min_n: int = 1  # Lower bound for the interval
     max_n: int = 10_000  # Upper bound for the interval
 
     size: int = 500  # Virtual dataset size
@@ -25,7 +27,8 @@ class CountPrimesConfig:
 
     def validate(self):
         """Validate configuration parameters"""
-        assert 1 <= self.max_n, "max_n must be at least 1"
+        assert 1 <= self.min_n, "min_n must be at least 1"
+        assert self.min_n <= self.max_n, "min_n must be less than or equal to max_n"
 
 
 class CountPrimesDataset(ProceduralDataset):
@@ -49,15 +52,42 @@ class CountPrimesDataset(ProceduralDataset):
     def __getitem__(self, idx: int) -> dict:
         """Generate a single Count Primes question"""
         rng = Random(self.seed + idx)
-        start = rng.randint(1, self.config.max_n)
+        start = rng.randint(self.config.min_n, self.config.max_n)
         end = rng.randint(start, self.config.max_n)
         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),
-            "metadata": {"start": start, "end": end, "primes": primes, "solution": answer},
+            "metadata": {
+                "start": start,
+                "end": end,
+                "primes": primes,
+                "solution": answer,
+                "difficulty": {
+                    "n": (start, end),
+                },
+            },
         }
 
 
-register_dataset("count_primes", CountPrimesDataset, CountPrimesConfig)
+class CountPrimesCurriculum(BaseCurriculum):
+    def __init__(self):
+        super().__init__(CountPrimesCurriculum.__name__, CountPrimesConfig)
+
+        # Define attributes
+        self._define_attributes(
+            RangeAttributeDefinition(
+                name="n",
+                levels=[1000, 10_000, 50_000, 100_000],
+                default_level=0,
+                description="Up to which number to consider the primes",
+                attr_type=AttributeType.APPEND,
+                min_value=1,
+                lower_field_name="min_n",
+                upper_field_name="max_n",
+            )
+        )
+
+
+register_dataset("count_primes", CountPrimesDataset, CountPrimesConfig, CountPrimesCurriculum)