include ranges rather than sampled values in difficulty metadata dicts (#387)

* update difficulty metadata for logic datasets

* update difficulty metadata for graph datasets

* update difficulty metadata for geometry datasets

* update difficulty metadata for games datasets

* update difficulty metadata for cognition datasets

* update difficulty metadata for arithmetic datasets

* update difficulty metadata for arc datasets

* update difficulty metadata for algorithmic datasets

* update difficulty metadata for algebra datasets

* use tuples

* update tests

* update tests

reasoning_gym/cognition/color_cube_rotation.py

@@ -141,7 +141,9 @@ class ColorCubeRotationDataset(ProceduralDataset):
                 "rotations": [r.value for r in rotations],
                 "target_side": target_side.value,
                 "num_rotations": num_rotations,
-                "difficulty": {"rotations": num_rotations},
+                "difficulty": {
+                    "rotations": (self.config.min_rotations, self.config.max_rotations),
+                },
             },
         }
 

reasoning_gym/cognition/figlet_fonts.py

@@ -188,7 +188,9 @@ class FigletFontDataset(ProceduralDataset):
             "metadata": {
                 "font": chosen_font,
                 "space_letters": self.config.space_letters,
-                "difficulty": {"word_len": len(word)},
+                "difficulty": {
+                    "word_len": (self.config.min_word_len, self.config.max_word_len),
+                },
             },
         }
 

reasoning_gym/cognition/modulo_grid.py

@@ -140,9 +140,11 @@ class ModuloGridDataset(ProceduralDataset):
                 "target": target,
                 "operation": operation,
                 "difficulty": {
-                    "holes": self.config.max_holes,
                     "size_x": self.config.size_x,
                     "size_y": self.config.size_y,
+                    "holes": self.config.max_holes,
+                    "divisor": self.config.max_divisor,
+                    "target": self.config.max_target,
                 },
             },
         }

reasoning_gym/cognition/needle_haystack.py

@@ -103,7 +103,13 @@ class NeedleHaystackDataset(ProceduralDataset):
         return {
             "question": full_text,
             "answer": stack["needle"][0],
-            "metadata": {"question": question, "difficulty": {"num_statements": num_statements}},
+            "metadata": {
+                "question": question,
+                "num_statements": num_statements,
+                "difficulty": {
+                    "num_statements": (self.config.min_num_statements, self.config.max_num_statements),
+                },
+            },
         }
 
     def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float:

reasoning_gym/cognition/number_sequences.py

@@ -195,7 +195,14 @@ class NumberSequenceDataset(ProceduralDataset):
         return {
             "question": ", ".join(map(str, visible_terms)) + ", ?",
             "answer": str(sequence[-1]),
-            "metadata": {"rule": rule.to_string(), "complexity": complexity, "sequence": sequence},
+            "metadata": {
+                "rule": rule.to_string(),
+                "complexity": complexity,
+                "sequence": sequence,
+                "difficulty": {
+                    "max_complexity": self.config.max_complexity,
+                },
+            },
         }
 
 

reasoning_gym/cognition/rectangle_count.py

@@ -117,7 +117,14 @@ class RectangleCountDataset(ProceduralDataset):
         return {
             "question": QUESTION_TEMPLATE.format(puzzle=puzzle),
             "answer": str(answer),
-            "metadata": {"puzzle": puzzle, "solution": answer, "difficulty": {"max_rectangles": target}},
+            "metadata": {
+                "puzzle": puzzle,
+                "solution": answer,
+                "num_rectangles": target,
+                "difficulty": {
+                    "max_rectangles": self.config.max_rectangles,
+                },
+            },
         }
 
     def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float:

reasoning_gym/cognition/rubiks_cube.py

@@ -110,8 +110,8 @@ class RubiksCubeDataset(ProceduralDataset):
                 "scramble_moves": " ".join([str(move) for move in scramble_moves]),
                 "example_correct_answer": actions_string,
                 "difficulty": {
-                    "scramble_steps": num_steps,
                     "cube_size": self.config.cube_size,
+                    "scramble_steps": (self.config.min_scramble_steps, self.config.max_scramble_steps),
                 },
             },
         }