fix(dataset): use conservative bounds for quantile aggregation instead of incorrect weighted mean (#3804)

* fix(stats): use conservative bounds for quantile aggregation instead of incorrect weighted mean

* docs: add --overwrite/--skip-images/--root options to augment_dataset_quantile_stats usage

* fix(dataset): clarify quantile aggregation semantics

* fix(augment): handle quantile stats edge cases

tests/datasets/test_augment_quantile_stats.py

@@ -12,8 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from types import SimpleNamespace
+
 import numpy as np
 import pytest
+import torch
 
 pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
 
@@ -24,7 +27,9 @@ from lerobot.scripts.augment_dataset_quantile_stats import (
 
 
 def _numeric_keys(dataset):
-    return [k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string")]
+    return [
+        k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string", "language")
+    ]
 
 
 def _image_keys(dataset):
@@ -102,3 +107,112 @@ def test_quantile_stats_present_after_compute(tmp_path, lerobot_dataset_factory)
     )
     stats = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
     assert has_quantile_stats(stats)
+
+
+class FakeHFDataset:
+    """Minimal stand-in exposing the column slicing used by the augment script."""
+
+    def __init__(self, columns: dict[str, list]):
+        self._columns = columns
+
+    def select_columns(self, keys):
+        return FakeHFDataset({key: self._columns[key] for key in keys})
+
+    def __getitem__(self, index):
+        return {key: values[index] for key, values in self._columns.items()}
+
+
+def test_compute_quantile_stats_skips_language_features():
+    class FakeDataset:
+        num_episodes = 1
+        features = {
+            "action": {"dtype": "float32"},
+            "observation.language": {"dtype": "language"},
+        }
+        meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 2}])
+        hf_dataset = FakeHFDataset(
+            {
+                "action": [[0.0], [1.0]],
+                "observation.language": [
+                    [{"role": "user", "content": "pick"}],
+                    [{"role": "assistant", "content": "done"}],
+                ],
+            }
+        )
+
+    stats = compute_quantile_stats_for_dataset(FakeDataset())
+
+    assert set(stats) == {"action"}
+
+
+def test_compute_quantile_stats_skip_images_avoids_decoding():
+    class FakeDataset:
+        num_episodes = 1
+        features = {
+            "action": {"dtype": "float32"},
+            "observation.images.cam": {"dtype": "video"},
+        }
+        meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 2}])
+        hf_dataset = FakeHFDataset({"action": [[0.0], [1.0]]})
+
+        def __getitem__(self, index):
+            raise AssertionError(f"video frame {index} was decoded despite skip_images=True")
+
+    stats = compute_quantile_stats_for_dataset(FakeDataset(), skip_images=True)
+
+    assert set(stats) == {"action"}
+
+
+def test_compute_quantile_stats_handles_single_frame():
+    class FakeDataset:
+        num_episodes = 1
+        features = {"action": {"dtype": "float32"}}
+        meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 1}])
+        hf_dataset = FakeHFDataset({"action": [[5.0, 7.0]]})
+
+    stats = compute_quantile_stats_for_dataset(FakeDataset())
+
+    np.testing.assert_array_equal(stats["action"]["count"], np.array([1]))
+    for key in ("min", "max", "mean", "q01", "q10", "q50", "q90", "q99"):
+        np.testing.assert_allclose(stats["action"][key], np.array([5.0, 7.0]))
+
+
+def test_compute_quantile_stats_image_count_uses_frames():
+    frames = [torch.zeros(3, 2, 2), torch.ones(3, 2, 2)]
+
+    class FakeDataset:
+        num_episodes = 1
+        features = {"observation.images.cam": {"dtype": "video"}}
+        meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 2}])
+        hf_dataset = FakeHFDataset({})
+
+        def __getitem__(self, index):
+            return {"observation.images.cam": frames[index]}
+
+    stats = compute_quantile_stats_for_dataset(FakeDataset(), use_sampling=False)
+    image_stats = stats["observation.images.cam"]
+
+    np.testing.assert_array_equal(image_stats["count"], np.array([2]))
+    assert image_stats["mean"].shape == (3, 1, 1)
+    np.testing.assert_allclose(image_stats["mean"], np.full((3, 1, 1), 0.5))
+
+
+def test_compute_quantile_stats_accumulates_across_episodes():
+    values = [[float(value)] for value in range(100)] + [[float(value)] for value in range(1000, 1010)]
+
+    class FakeDataset:
+        num_episodes = 2
+        features = {"action": {"dtype": "float32"}}
+        meta = SimpleNamespace(
+            episodes=[
+                {"dataset_from_index": 0, "dataset_to_index": 100},
+                {"dataset_from_index": 100, "dataset_to_index": 110},
+            ]
+        )
+        hf_dataset = FakeHFDataset({"action": values})
+
+    stats = compute_quantile_stats_for_dataset(FakeDataset())
+
+    np.testing.assert_array_equal(stats["action"]["count"], np.array([110]))
+    expected_q90 = np.percentile(np.asarray(values), 90, axis=0)
+    np.testing.assert_allclose(stats["action"]["q90"], expected_q90, atol=0.1)