Merge remote-tracking branch 'origin/main' into feat/curr-adj

reasoning_gym/arc/__init__.py

@@ -1,10 +1,11 @@
-from .arc_1d import Arc1DConfig, Arc1DDataset
+from .arc_1d import Arc1DConfig, Arc1DCurriculum, Arc1DDataset
 from .arc_agi import ArcAgiConfig, ArcAgiDataset
 from .rearc import ReArcConfig, ReArcCurriculum, ReArcDataset
 
 __all__ = [
     "Arc1DConfig",
     "Arc1DDataset",
+    "Arc1DCurriculum",
     "ArcAgiConfig",
     "ArcAgiDataset",
     "ReArcDataset",

reasoning_gym/arc/arc_1d.py

@@ -2,6 +2,7 @@ from dataclasses import dataclass
 from random import Random
 from typing import Optional
 
+from ..coaching import BaseCurriculum, RangeAttributeDefinition
 from ..dataset import ProceduralDataset
 from ..factory import register_dataset
 
@@ -108,9 +109,31 @@ class Arc1DDataset(ProceduralDataset):
                 "size": size,
                 "train_examples": train_examples,
                 "test_example": test_example,
+                "difficulty": {
+                    "size": (self.config.min_size, self.config.max_size),
+                },
             },
         }
 
 
+class Arc1DCurriculum(BaseCurriculum):
+    """Curriculum for ARC 1D tasks"""
+
+    def __init__(self):
+        super().__init__(Arc1DCurriculum.__name__, Arc1DConfig)
+
+        # Define attributes
+        self._define_attributes(
+            RangeAttributeDefinition(
+                name="size",
+                levels=[10, 25, 50, 100],
+                lower_field_name="min_size",
+                upper_field_name="max_size",
+                description="Grid size",
+                ensure_interval=True,
+            )
+        )
+
+
 # Register the dataset
-register_dataset(DATASET_NAME, Arc1DDataset, Arc1DConfig)
+register_dataset(DATASET_NAME, Arc1DDataset, Arc1DConfig, Arc1DCurriculum)

reasoning_gym/logic/__init__.py

@@ -4,7 +4,7 @@ Logic tasks for training reasoning capabilities.
 
 from .aiw import AliceInWonderlandConfig, AliceInWonderlandCurriculum, AliceInWonderlandDataset
 from .circuit_logic import CircuitLogicConfig, CircuitLogicCurriculum, CircuitLogicDataset
-from .knights_knaves import KnightsKnavesConfig, KnightsKnavesDataset
+from .knights_knaves import KnightsKnavesConfig, KnightsKnavesCurriculum, KnightsKnavesDataset
 from .propositional_logic import PropositionalLogicConfig, PropositionalLogicCurriculum, PropositionalLogicDataset
 from .self_reference import SelfReferenceConfig, SelfReferenceCurriculum, SelfReferenceDataset
 from .syllogisms import SyllogismConfig, SyllogismDataset
@@ -31,4 +31,5 @@ __all__ = [
     "CircuitLogicCurriculum",
     "KnightsKnavesConfig",
     "KnightsKnavesDataset",
+    "KnightsKnavesCurriculum",
 ]

reasoning_gym/logic/knights_knaves.py

@@ -8,6 +8,8 @@ import numpy as np
 
 from reasoning_gym.factory import ProceduralDataset, register_dataset
 
+from ..coaching import BaseCurriculum, ScalarAttributeDefinition
+
 DATASET_NAME = "knights_knaves"
 
 COMMON_NAMES = [
@@ -462,6 +464,11 @@ class KnightsKnavesDataset(ProceduralDataset):
             "solution": problem["solution"],
             "names": formatted["names"],
             "knight_knave_terms": formatted["knight_knave"],
+            "difficulty": {
+                "n_people": self.config.n_people,
+                "depth_constraint": self.config.depth_constraint,
+                "width_constraint": self.config.width_constraint,
+            },
         }
 
         return {"question": question, "answer": answer, "metadata": metadata}
@@ -515,4 +522,30 @@ class KnightsKnavesDataset(ProceduralDataset):
         return 0.0
 
 
-register_dataset(DATASET_NAME, KnightsKnavesDataset, KnightsKnavesConfig)
+class KnightsKnavesCurriculum(BaseCurriculum):
+    def __init__(self):
+        super().__init__(KnightsKnavesCurriculum.__name__, KnightsKnavesConfig)
+
+        self._define_attributes(
+            ScalarAttributeDefinition(
+                name="n_people",
+                levels=[2, 3, 4, 5],
+                description="Number of people in the problem",
+                field_name="n_people",
+            ),
+            ScalarAttributeDefinition(
+                name="depth_constraint",
+                levels=[2, 3, 4, 5],
+                description="Depth of the problem",
+                field_name="depth_constraint",
+            ),
+            ScalarAttributeDefinition(
+                name="width_constraint",
+                levels=[2, 3, 4, 5],
+                description="Width of the problem",
+                field_name="width_constraint",
+            ),
+        )
+
+
+register_dataset(DATASET_NAME, KnightsKnavesDataset, KnightsKnavesConfig, KnightsKnavesCurriculum)