feat(env): Count Bits Curriculum (#267)

* add min n

* count bits

reasoning_gym/arithmetic/__init__.py

@@ -6,7 +6,7 @@ from .basic_arithmetic import BasicArithmeticDataset, BasicArithmeticDatasetConf
 from .bitwise_arithmetic import BitwiseArithmeticConfig, BitwiseArithmeticDataset
 from .calendar_arithmetic import CalendarArithmeticConfig, CalendarArithmeticDataset
 from .chain_sum import ChainSumConfig, ChainSumDataset
-from .count_bits import CountBitsConfig, CountBitsDataset
+from .count_bits import CountBitsConfig, CountBitsCurriculum, CountBitsDataset
 from .decimal_arithmetic import DecimalArithmeticConfig, DecimalArithmeticDataset
 from .decimal_chain_sum import DecimalChainSumConfig, DecimalChainSumDataset
 from .dice import DiceConfig, DiceDataset
@@ -48,6 +48,7 @@ __all__ = [
     "TimeIntervalsDataset",
     "CountBitsConfig",
     "CountBitsDataset",
+    "CountBitsCurriculum",
     "DiceConfig",
     "DiceDataset",
     "NumberFormatConfig",

reasoning_gym/arithmetic/count_bits.py

@@ -4,6 +4,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 = """How many 1 bits are there in the binary representation of the number {number}?"""
@@ -13,6 +14,7 @@ QUESTION_TEMPLATE = """How many 1 bits are there in the binary representation of
 class CountBitsConfig:
     """Configuration for Count Bits dataset generation"""
 
+    min_n: int = 1  # Minimum number to consider
     max_n: int = 2**31 - 1  # Maximum number to consider
 
     size: int = 500  # Virtual dataset size
@@ -20,7 +22,7 @@ class CountBitsConfig:
 
     def validate(self):
         """Validate configuration parameters"""
-        assert 1 <= self.max_n, "max_n must be at least 1"
+        assert 1 <= self.min_n <= self.max_n, "min_n must be between 1 and max_n"
 
 
 class CountBitsDataset(ProceduralDataset):
@@ -33,7 +35,7 @@ class CountBitsDataset(ProceduralDataset):
         """Generate a single Count Bits question"""
         rng = Random(self.seed + idx)
 
-        number = rng.randint(1, self.config.max_n)
+        number = rng.randint(self.config.min_n, self.config.max_n)
         binary = bin(number)[2:]
         answer = binary.count("1")
 
@@ -44,4 +46,23 @@ class CountBitsDataset(ProceduralDataset):
         }
 
 
+class CountBitsCurriculum(BaseCurriculum):
+    def __init__(self):
+        super().__init__(CountBitsCurriculum.__name__, CountBitsConfig)
+
+        # Define attributes
+        self._define_attributes(
+            RangeAttributeDefinition(
+                name="n",
+                levels=[1_000, 1_000_000, 100_000_000, 2**31 - 1],
+                default_level=0,
+                description="Number to count bits in",
+                attr_type=AttributeType.APPEND,
+                min_value=1,
+                lower_field_name="min_n",
+                upper_field_name="max_n",
+            ),
+        )
+
+
 register_dataset("count_bits", CountBitsDataset, CountBitsConfig)