diff --git a/reasoning_gym/algorithmic/caesar_cipher.py b/reasoning_gym/algorithmic/caesar_cipher.py
index 0f628e39..6a69159a 100644
--- a/reasoning_gym/algorithmic/caesar_cipher.py
+++ b/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,
+                },
             },
         }
 
diff --git a/reasoning_gym/algorithmic/graph_color.py b/reasoning_gym/algorithmic/graph_color.py
index d1e17bf0..408a9f43 100644
--- a/reasoning_gym/algorithmic/graph_color.py
+++ b/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:
diff --git a/reasoning_gym/arithmetic/decimal_arithmetic.py b/reasoning_gym/arithmetic/decimal_arithmetic.py
index 2ebaf26c..17b54765 100644
--- a/reasoning_gym/arithmetic/decimal_arithmetic.py
+++ b/reasoning_gym/arithmetic/decimal_arithmetic.py
@@ -188,7 +188,12 @@ class DecimalArithmeticDataset(ProceduralDataset):
         return {
             "question": problem_str,
             "answer": str(answer),
-            "metadata": {"decimal_places": decimal_places, "num_terms": terms},
+            "metadata": {
+                "difficulty": {
+                    "decimal_places": decimal_places,
+                    "num_terms": terms,
+                },
+            },
         }
 
     def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float: