create init and add difficulty metadata for curriculum

Author: kumaranant1

reasoning_gym/probability/__init__.py

@@ -0,0 +1,11 @@
+"""
+Probability reasoning tasks.
+"""
+
+from .coin_flip import CoinFlipDataset, CoinFlipConfig, CoinFlipCurriculum
+
+__all__ = [
+    "CoinFlipDataset",
+    "CoinFlipConfig", 
+    "CoinFlipCurriculum"
+]

reasoning_gym/probability/coin_flip.py

@@ -25,6 +25,7 @@ def validate(self) -> None:
         assert self.size > 0, "size must be positive"
         assert self.min_trials > 0, "min_trials must be positive"
         assert self.max_trials >= self.min_trials, "max_trials must be >= min_trials"
+        assert self.allow_exact or self.allow_at_least, "At least one of allow_exact or allow_at_least must be True"
 
 
 class CoinFlipDataset(ProceduralDataset):
@@ -55,9 +56,6 @@ def __getitem__(self, idx: int) -> dict:
             available_types.append("exact")
         if self.config.allow_at_least:
             available_types.append("at_least")
-
-        if not available_types:
-            available_types = ["exact"]
 
         problem_type = rng.choice(available_types)
 
@@ -79,13 +77,16 @@ def __getitem__(self, idx: int) -> dict:
             "metadata": {
                 "source_dataset": DATASET_NAME,
                 "source_index": idx,
-                "num_tosses": n,
+                "num_trials": n,
                 "k_heads": k,
                 "problem_type": problem_type,
                 "rational": {
                     "numerator": self._rational_numerator(n, k, problem_type),
                     "denominator": 2 ** n,
-                }
+                },
+                "difficulty": {
+                    "num_trials": (self.config.min_trials, self.config.max_trials),
+                },
             }
         }
 

tests/test_coin_flip.py

@@ -1,6 +1,6 @@
 import pytest
 from fractions import Fraction
-from reasoning_gym.probability.coin_flip import CoinFlipDataset, CoinFlipConfig, CoinFlipCurriculum
+from reasoning_gym.probability import CoinFlipDataset, CoinFlipConfig, CoinFlipCurriculum
 
 def test_coin_flip_config_validation():
     """Test that invalid configs raise errors"""
@@ -16,6 +16,9 @@ def test_coin_flip_config_validation():
         config = CoinFlipConfig(min_trials=5, max_trials=3)
         config.validate()
 
+    with pytest.raises(AssertionError):
+        config = CoinFlipConfig(allow_exact=False, allow_at_least=False)
+        config.validate()
 
 def test_coin_flip_deterministic():
     """Dataset generates same items with same seed"""
@@ -41,13 +44,13 @@ def test_coin_flip_items():
 
 
         metadata = item["metadata"]
-        assert "num_tosses" in metadata
+        assert "num_trials" in metadata
         assert "k_heads" in metadata
         assert "problem_type" in metadata
         assert metadata["problem_type"] in ["exact", "at_least"]
 
         rational = metadata["rational"]
-        assert rational["denominator"] == 2 ** metadata["num_tosses"]
+        assert rational["denominator"] == 2 ** metadata["num_trials"]
         assert rational["numerator"] > 0