add metadata for caesar cipher, graph coloring, decimal arithmetic (#304)

* add metadata for caesar cipher, graph coloring, decimal arithmetic

* delete comma

* clean up variables

reasoning_gym/algorithmic/caesar_cipher.py

@@ -67,6 +67,7 @@ class CaesarCipherDataset(ProceduralDataset):
 
         # Select random sentence and rotation
         sentence = rng.choice(self.valid_sentences)
+        num_words = len(sentence.split())
         rotation = rng.randint(self.config.min_rotation, self.config.max_rotation)
 
         # Generate cipher text
@@ -79,6 +80,10 @@ class CaesarCipherDataset(ProceduralDataset):
                 "rotation": rotation,
                 "cipher_text": cipher_text,
                 "clear_text": sentence,
+                "difficulty": {
+                    "rotation": rotation,
+                    "words": num_words,
+                },
             },
         }
 

reasoning_gym/algorithmic/graph_color.py

@@ -187,12 +187,14 @@ class GraphColorDataset(ProceduralDataset):
 
         puzzle = None
         solution = None
+        num_vertices = rng.randint(self.config.min_num_vertices, self.config.max_num_vertices)
+        num_colors = rng.randint(self.config.min_num_colors, self.config.max_num_colors)
         while solution is None:
             puzzle = generate_graph_coloring_puzzle(
                 rng=rng,
-                num_vertices=rng.randint(self.config.min_num_vertices, self.config.max_num_vertices),
+                num_vertices=num_vertices,
                 edge_probability=self.config.edge_probability,
-                num_colors=rng.randint(self.config.min_num_colors, self.config.max_num_colors),
+                num_colors=num_colors,
             )
             solution = greedy_graph_coloring(puzzle)
 
@@ -209,7 +211,11 @@ Return your solution as a JSON map of vertices to colors. (For example: {{"0": 1
         return {
             "question": question,
             "answer": None,
-            "metadata": {"possible_answer": solution, "puzzle": puzzle},
+            "metadata": {
+                "possible_answer": solution,
+                "puzzle": puzzle,
+                "difficulty": {"num_vertices": num_vertices, "num_colors": num_colors},
+            },
         }
 
     def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float: