Fix aggregation rounding in ng_prepare_data

nemo_gym/train_data_utils.py

@@ -142,19 +142,28 @@ def _add(self: Self, other: Self) -> None:
         self.tdigest = self.tdigest + other.tdigest
 
     def _aggregate(self: Self) -> Self:
+        def round_metric(x: float) -> float:
+            if x >= 1 or x <= -1:
+                return round(x, 2)
+            return round(x, 3)
+
         n = self.total
         mean = self.mean if n > 0 else 0.0
         stddev = sqrt(self.M2 / (n - 1)) if n > 1 else 0.0
         med = float(self.tdigest.percentile(50)) if n > 0 and self.tdigest.n > 0 else 0.0
 
-        return AvgMinMax(
-            total=self.total,
-            average=mean,
-            min=self.min if n > 0 else 0.0,
-            max=self.max if n > 0 else 0.0,
-            median=med,
-            stddev=stddev,
-        )
+        params = {
+            "total": self.total,
+            "average": mean,
+            "min": self.min if n > 0 else 0.0,
+            "max": self.max if n > 0 else 0.0,
+            "median": med,
+            "stddev": stddev,
+        }
+
+        final_params = {k: round_metric(v) if isinstance(v, float) else v for k, v in params.items()}
+
+        return AvgMinMax(**final_params)
 
 
 class StringMetrics(BaseModel):

resources_servers/comp_coding/data/example_metrics.json

@@ -19,7 +19,7 @@
         "Min": 348.0,
         "Max": 542.0,
         "Median": 473.0,
-        "Standard deviation": 79.75587752636166
+        "Standard deviation": 79.76
     },
     "Number of turns": {
         "Total # non-null values": 5,

resources_servers/equivalence_llm_judge/data/example_metrics.json

@@ -19,7 +19,7 @@
         "Min": 29.0,
         "Max": 44.0,
         "Median": 32.0,
-        "Standard deviation": 6.308724118235001
+        "Standard deviation": 6.31
     },
     "Number of turns": {
         "Total # non-null values": 5,

resources_servers/google_search/data/example_metrics.json

@@ -19,7 +19,7 @@
         "Min": 230.0,
         "Max": 333.0,
         "Median": 293.0,
-        "Standard deviation": 40.44996909763962
+        "Standard deviation": 40.45
     },
     "Number of turns": {
         "Total # non-null values": 5,
@@ -37,16 +37,16 @@
         "Median": 0.6,
         "Standard deviation": 0.0
     },
+    "expected_answer": {
+        "unique_count": 3,
+        "total_count": 5
+    },
     "task_difficulty_qwen3_32b_avg_8": {
         "Total # non-null values": 5,
         "Average": 0.4,
         "Min": 0.25,
         "Max": 0.5,
-        "Median": 0.40625,
-        "Standard deviation": 0.10458250331675945
-    },
-    "expected_answer": {
-        "unique_count": 3,
-        "total_count": 5
+        "Median": 0.406,
+        "Standard deviation": 0.105
     }
 }

resources_servers/instruction_following/data/example_metrics.json

@@ -19,7 +19,7 @@
         "Min": 57.0,
         "Max": 180.0,
         "Median": 74.0,
-        "Standard deviation": 49.95297788921097
+        "Standard deviation": 49.95
     },
     "Number of turns": {
         "Total # non-null values": 5,
@@ -43,7 +43,7 @@
         "Min": 2228.0,
         "Max": 44655.0,
         "Median": 17616.0,
-        "Standard deviation": 20015.806548825356
+        "Standard deviation": 20015.81
     },
     "instruction_id_list": {
         "unique_count": 13,

resources_servers/library_judge_math/data/example_metrics.json

@@ -19,7 +19,7 @@
         "Min": 45.0,
         "Max": 120.0,
         "Median": 63.0,
-        "Standard deviation": 30.405591591021544
+        "Standard deviation": 30.41
     },
     "Number of turns": {
         "Total # non-null values": 5,

resources_servers/mcqa/data/example_metrics.json

@@ -19,7 +19,7 @@
         "Min": 71.0,
         "Max": 210.0,
         "Median": 178.0,
-        "Standard deviation": 66.37017402418047
+        "Standard deviation": 66.37
     },
     "Number of turns": {
         "Total # non-null values": 5,

resources_servers/multineedle/data/example_metrics.json

@@ -19,7 +19,7 @@
         "Min": 1499.0,
         "Max": 1509.0,
         "Median": 1505.0,
-        "Standard deviation": 4.000000000000012
+        "Standard deviation": 4.0
     },
     "Number of turns": {
         "Total # non-null values": 5,
@@ -43,15 +43,23 @@
         "Min": 0.0,
         "Max": 4.0,
         "Median": 2.0,
-        "Standard deviation": 1.5811388300841898
+        "Standard deviation": 1.58
+    },
+    "expected_synonyms": {
+        "unique_count": 2,
+        "total_count": 10
     },
     "expected_synonym_values": {
         "Total # non-null values": 10,
         "Average": 559.0,
         "Min": 407.0,
         "Max": 711.0,
         "Median": 559.0,
-        "Standard deviation": 160.22206811519789
+        "Standard deviation": 160.22
+    },
+    "minefield_label": {
+        "unique_count": 1,
+        "total_count": 5
     },
     "minefield_label_value": {
         "Total # non-null values": 5,
@@ -60,13 +68,5 @@
         "Max": 299.0,
         "Median": 299.0,
         "Standard deviation": 0.0
-    },
-    "expected_synonyms": {
-        "unique_count": 2,
-        "total_count": 10
-    },
-    "minefield_label": {
-        "unique_count": 1,
-        "total_count": 5
     }
 }

resources_servers/multiverse_math_hard/data/example_metrics.json

@@ -19,7 +19,7 @@
         "Min": 402.0,
         "Max": 429.0,
         "Median": 408.0,
-        "Standard deviation": 10.784247771634329
+        "Standard deviation": 10.78
     },
     "Number of turns": {
         "Total # non-null values": 5,
@@ -43,26 +43,26 @@
         "Min": 0.0,
         "Max": 4.0,
         "Median": 2.0,
-        "Standard deviation": 1.5811388300841898
+        "Standard deviation": 1.58
+    },
+    "ground_truth": {
+        "unique_count": 5,
+        "total_count": 5
     },
     "depth": {
         "Total # non-null values": 5,
         "Average": 2.4,
         "Min": 1.0,
         "Max": 4.0,
         "Median": 2.0,
-        "Standard deviation": 1.51657508881031
+        "Standard deviation": 1.52
     },
     "breadth": {
         "Total # non-null values": 5,
         "Average": 3.0,
         "Min": 1.0,
         "Max": 5.0,
         "Median": 3.25,
-        "Standard deviation": 1.8708286933869707
-    },
-    "ground_truth": {
-        "unique_count": 5,
-        "total_count": 5
+        "Standard deviation": 1.87
     }
 }

resources_servers/python_math_exec/data/example_metrics.json

@@ -19,7 +19,7 @@
         "Min": 190.0,
         "Max": 265.0,
         "Median": 208.0,
-        "Standard deviation": 30.405591591021544
+        "Standard deviation": 30.41
     },
     "Number of turns": {
         "Total # non-null values": 5,
@@ -43,7 +43,7 @@
         "Min": 0.0,
         "Max": 4.0,
         "Median": 2.0,
-        "Standard deviation": 1.5811388300841898
+        "Standard deviation": 1.58
     },
     "expected_result": {
         "unique_count": 5,

resources_servers/simple_weather/data/example_metrics.json

@@ -19,7 +19,7 @@
         "Min": 51.0,
         "Max": 57.0,
         "Median": 52.5,
-        "Standard deviation": 2.4494897427831783
+        "Standard deviation": 2.45
     },
     "Number of turns": {
         "Total # non-null values": 5,

resources_servers/stateful_counter/data/example_metrics.json

@@ -43,14 +43,14 @@
         "Min": 6.0,
         "Max": 42.0,
         "Median": 24.0,
-        "Standard deviation": 14.230249470757707
+        "Standard deviation": 14.23
     },
     "initial_count": {
         "Total # non-null values": 5,
         "Average": 9.0,
         "Min": 3.0,
         "Max": 15.0,
         "Median": 9.0,
-        "Standard deviation": 4.743416490252569
+        "Standard deviation": 4.74
     }
 }

resources_servers/workbench/data/example_metrics.json

@@ -19,7 +19,7 @@
         "Min": 1726.0,
         "Max": 1737.0,
         "Median": 1731.0,
-        "Standard deviation": 4.764451699828638
+        "Standard deviation": 4.76
     },
     "Number of turns": {
         "Total # non-null values": 5,
@@ -43,7 +43,7 @@
         "Min": 0.0,
         "Max": 4.0,
         "Median": 2.0,
-        "Standard deviation": 1.5811388300841898
+        "Standard deviation": 1.58
     },
     "category": {
         "unique_count": 3,

tests/unit_tests/test_dataset_viewer.py

@@ -92,14 +92,14 @@ def mock_compute_sample_metrics(line: str):
         # Check computed values
         reward_stats = result_1["reward"]
         assert reward_stats["Total # non-null values"] == 3
-        assert reward_stats["Average"] == (1.0 + 0.0 + 0.5) / 3
+        assert reward_stats["Average"] == round(((1.0 + 0.0 + 0.5) / 3), 3)
         assert reward_stats["Min"] == 0.0
         assert reward_stats["Max"] == 1.0
 
         # Check computed values with bools converted to int
         accuracy_stats = result_1["accuracy"]
         assert accuracy_stats["Total # non-null values"] == 3
-        assert accuracy_stats["Average"] == (1 + 0 + 1) / 3
+        assert accuracy_stats["Average"] == round(((1 + 0 + 1) / 3), 3)
         assert accuracy_stats["Min"] == 0
         assert accuracy_stats["Max"] == 1
 

tests/unit_tests/test_train_data_utils.py

@@ -353,7 +353,7 @@ def custom_open(filename, mode="r"):
                     min=0.0,
                     max=4.0,
                     median=2.0,
-                    stddev=1.5811388300841898,
+                    stddev=1.58,
                 ),
                 expected_synonym_values=AvgMinMax(
                     is_aggregated=True,
@@ -362,7 +362,7 @@ def custom_open(filename, mode="r"):
                     min=407.0,
                     max=711.0,
                     median=559.0,
-                    stddev=160.22206811519789,
+                    stddev=160.22,
                 ),
                 minefield_label_value=AvgMinMax(
                     is_aggregated=True,