Add OpenPi, Pi0 and Pi0.5 (#1910)

* initial commit * change device in test * do detailed import * adhere to python 3.11 syntax * fix autodocstring * additionally * do same in other files * add model. prefix to all keys in state dict * use dummy stats * add pi05 * also shorten action_steps * fix test * all test pass! and fix tokenizer max length between 05 and 0 * remove test * fix transformer dependency * fix test * split pi0 and pi05 policy in seperate files * fix test * fix push to hub test * add some comments, license and readme * remove warning in config * add pi05 to factory * remove check * rename action_horizon to chunk_size * clean up padding of state and action (more in line with lerobot pi0) * add openpi image transforms for training and add more flexibility to _preprocess_images similar to lerobot pi0 * fix key match from pytorch state dict (similar keys to openpi implementation now) * also for pi05 * update to python 3.11 * revert to openpi transformer replace python 3.11 * fix(modeling pi0): nit warning message * use safeauto_docstring * fix: remove unused param * fix from pretrained * add preprocess tests * also compile forward method * Do not add model prefix to normalization * use same name for action and state dim as lerobot pi0 and remove fixed image keys * load from pretrained_path * temp: hardcode base model * fix override self.pretrained_path = None overwrite * rename to loss * remove additional image augmentations, lerobot dataset already does this * Add docs * put tests in test folder * Add test to instatiate all base models * go back to python 3.10 * update docs * adapt docs pi05 * change docs: finetune base model options * minor docs fixes and dependencies * remove todo * cast float64 to float32 for mps * skip if no transformers * fix tests * add new models to modelcard * add back init * fix circular input * feat: only run pi test on GPU * remove require_nightly_gpu * replace decorator test_pi0_openpi * rename action_dim, state_dim to max_action_dim, max_state_dim * fix doc and constants * cleanup tests * fix from pretrained * fix tests * add comment pi0 pi05 tests, add image features to pi0 pi05 hub tests * fix, state is included in language not in flow head * Move test to specific folder * and paligemma task with newline * remove add_special_tokens, not needed * feedback pr * Remove previous pi0 and rename pi0_openpi and pi05_openpi * Add Quantile stats to LeRobotDataset (#1985) * - Add RunningQuantileStats class for efficient histogram-based quantile computation - Integrate quantile parameters (compute_quantiles, quantiles) into LeRobotDataset - Support quantile computation during episode collection and aggregation - Add comprehensive function-based test suite (24 tests) for quantile functionality - Maintain full backward compatibility with existing stats computation - Enable configurable quantiles (default: [0.01, 0.99]) for robust normalization * style fixes, make quantiles computation by default to new datasets * fix tests * - Added DEFAULT_QUANTILES=[0.01, 0.10, 0.50, 0.90, 0.99] to be computed for each features instead of being chosen by the user - Fortified tests. * - add helper functions to reshape stats - add missing test for quantiles * - Add QUANTILE normalization mode to normalize the data with the 1st and 99th percentiles. - Add QUANTILE10 normalization mode to normalize the data with the 10th and 90th percentiles. * style fixes * Added missing lisence * Simplify compute_stats * - added script `augment_dataset_quantile_stats.py` so that we can add quantile stats to existing v3 datasets that dont have quatniles - modified quantile computation instead of using the edge for the value, interpolate the values in the bin * rename pi0/pi05 files * Remove open pi patch and use custom transformer branch for now * renaming * fix * Revert "fix" This reverts commit 1ea65730ac2cbca6e5869df734fbd4392561b3c6. * fix naming * feet(pi0/pi0.5): add pipeline (#2009) * feat(processor): convert openpi model with processor * TODO: Make test works * fix(modeling_pi0openpi): update attention mask value and time scaling; improve task handling in tests - Changed the attention mask value from `self.config.attention_mask_value` to a fixed value of `-2.3819763e38`. - Updated time scaling in the `sample_noise` method to use a constant factor of `0.999` and an offset of `0.001`. - Enhanced task handling in tests to ensure proper formatting and batch size consistency. - Cleaned up commented-out test code for clarity. * refactor(pi0): rename PI0OpenPIConfig and PI0OpenPIPolicy to PI0Config and PI0Policy - Updated imports and references throughout the codebase to reflect the new naming convention. - Introduced a new processor file for PI0 to handle pre-processing and post-processing steps. - Adjusted tests to utilize the renamed classes, ensuring consistency and functionality. - Enhanced clarity and maintainability by removing outdated naming conventions. * refactor(pi05): rename PI0OpenPIPolicy to PI0Policy and update configuration - Renamed `PI0OpenPIPolicy` to `PI0Policy` for consistency with naming conventions. - Updated the `PI05OpenPIConfig` to include a new `tokenizer_max_length` attribute and changed the normalization mode for state from `MEAN_STD` to `QUANTILES`. - Simplified model initialization in `PI05OpenPIPolicy` by removing unused `dataset_stats` parameter. - Added a new processor class for `Pi05PrepareStateTokenizerProcessorStep` with `@dataclass` for improved readability. - Introduced a test script to compare the integration of the PI0OpenPI policy with the original implementation, ensuring local testing compatibility. * feat(processor): convert openpi model with processor * TODO: Make test works * fix(modeling_pi0openpi): update attention mask value and time scaling; improve task handling in tests - Changed the attention mask value from `self.config.attention_mask_value` to a fixed value of `-2.3819763e38`. - Updated time scaling in the `sample_noise` method to use a constant factor of `0.999` and an offset of `0.001`. - Enhanced task handling in tests to ensure proper formatting and batch size consistency. - Cleaned up commented-out test code for clarity. * refactor(pi0): rename PI0OpenPIConfig and PI0OpenPIPolicy to PI0Config and PI0Policy - Updated imports and references throughout the codebase to reflect the new naming convention. - Introduced a new processor file for PI0 to handle pre-processing and post-processing steps. - Adjusted tests to utilize the renamed classes, ensuring consistency and functionality. - Enhanced clarity and maintainability by removing outdated naming conventions. * refactor(pi05): rename PI0OpenPIPolicy to PI0Policy and update configuration - Renamed `PI0OpenPIPolicy` to `PI0Policy` for consistency with naming conventions. - Updated the `PI05OpenPIConfig` to include a new `tokenizer_max_length` attribute and changed the normalization mode for state from `MEAN_STD` to `QUANTILES`. - Simplified model initialization in `PI05OpenPIPolicy` by removing unused `dataset_stats` parameter. - Added a new processor class for `Pi05PrepareStateTokenizerProcessorStep` with `@dataclass` for improved readability. - Introduced a test script to compare the integration of the PI0OpenPI policy with the original implementation, ensuring local testing compatibility. * refactor(pi05): update imports and rename configuration classes - Changed imports to reflect the new naming convention for PI05 configuration and policy classes. - Renamed `PI05OpenPIConfig` to `PI05Config` and `PI05OpenPIPolicy` to `PI05Policy` for consistency. - Introduced a new processor file for PI05, implementing pre-processing and post-processing steps. - Updated tests to utilize the renamed classes, ensuring functionality and consistency across the codebase. * update(pi05): increase tokenizer_max_length for improved processing - Changed the `tokenizer_max_length` from 48 to 200 to enhance the model's capability in handling longer sequences. - This adjustment aims to improve the overall performance and flexibility of the PI05 configuration. * add default for state (max_state_dim) * correct naming * fix import * cleanup code * remove unused test * us quantiles for action * move to device * remove discrete state assert * fix pi05 test * move pi05 to device * use base models in comparison tests * small renames for tests * change number of tokens pi05 test * fix openpi tokenization in test * fix hub test * fix test * assert lerobot vs openpi tests --------- Co-authored-by: Pepijn <pepijn@huggingface.co> * add headers * add back previously removed imports * update if statement load processor with dataset stats * remove to avoid circular import * inject dataset stats for pretrained models * check normalization before applying * add link to quantile augument script * fix(policies): transformers import for ci in PI0 & PI05 (#2039) * fix(policies): transformers import for ci in PI0 * fix(policies): transformers import for ci in PI05 * test(processor): fix expected raise when normalization types are missing (#2040) * switch normalization order pipeline for pi05 * Fix/quantiles script (#2064) * refactor augment stats with quantiles script add parallelization for faster processing shift the quantile normalization between -1 1 * fix replay buffer tests * fix comment * overwrite the pipeline normalization features with the policy features * remove double normalization overwrite * cleanup from pretrained * remove typo * also set norm_map * fix(augment_quantiles) images incorrectly divided by 255 * clamp quantiles * link to lerobot base models * rename tests * encorperate PR feedback * update docstring for RunningQuantileStats * update doc links * Revert "clamp quantiles" This reverts commit 172207471c8f2cb62958e9a9e6a0535ba3ff67d4. * fix self.paligemma * fix tests related to quantiles that were scaled to [0,1], the new range is [-1, 1] * fix libero doc and use different transformer branch * use fix branch instead of feat * update results libero * add new line * fix formatting * precommit * update results libero * update libero doc * update title * final changes * add quantiles to test * run pre commit --------- Signed-off-by: Steven Palma <imstevenpmwork@ieee.org> Co-authored-by: Michel Aractingi <michel.aractingi@huggingface.co> Co-authored-by: Adil Zouitine <adilzouitinegm@gmail.com> Co-authored-by: Steven Palma <imstevenpmwork@ieee.org> Co-authored-by: Steven Palma <steven.palma@huggingface.co>
2025-10-02 13:14:45 +02:00
parent b6c528a438
commit abde7be3b3
43 changed files with 5886 additions and 2288 deletions
--- a/tests/datasets/test_compute_stats.py
+++ b/tests/datasets/test_compute_stats.py
@@ -19,6 +19,7 @@ import numpy as np
 import pytest

 from lerobot.datasets.compute_stats import (
+    RunningQuantileStats,
    _assert_type_and_shape,
    aggregate_feature_stats,
    aggregate_stats,
@@ -102,6 +103,9 @@ def test_get_feature_stats_axis_1(sample_array):
        "count": np.array([3]),
    }
    result = get_feature_stats(sample_array, axis=(1,), keepdims=False)
+
+    # Check that basic stats are correct (quantiles are also included now)
+    assert set(expected.keys()).issubset(set(result.keys()))
    for key in expected:
        np.testing.assert_allclose(result[key], expected[key])

@@ -115,6 +119,9 @@ def test_get_feature_stats_no_axis(sample_array):
        "count": np.array([3]),
    }
    result = get_feature_stats(sample_array, axis=None, keepdims=False)
+
+    # Check that basic stats are correct (quantiles are also included now)
+    assert set(expected.keys()).issubset(set(result.keys()))
    for key in expected:
        np.testing.assert_allclose(result[key], expected[key])

@@ -308,3 +315,520 @@ def test_aggregate_stats():
            results[fkey]["std"], expected_agg_stats[fkey]["std"], atol=1e-04, rtol=1e-04
        )
        np.testing.assert_allclose(results[fkey]["count"], expected_agg_stats[fkey]["count"])
+
+
+def test_running_quantile_stats_initialization():
+    """Test proper initialization of RunningQuantileStats."""
+    running_stats = RunningQuantileStats()
+    assert running_stats._count == 0
+    assert running_stats._mean is None
+    assert running_stats._num_quantile_bins == 5000
+
+    # Test custom bin size
+    running_stats_custom = RunningQuantileStats(num_quantile_bins=1000)
+    assert running_stats_custom._num_quantile_bins == 1000
+
+
+def test_running_quantile_stats_single_batch_update():
+    """Test updating with a single batch."""
+    np.random.seed(42)
+    data = np.random.normal(0, 1, (100, 3))
+
+    running_stats = RunningQuantileStats()
+    running_stats.update(data)
+
+    assert running_stats._count == 100
+    assert running_stats._mean.shape == (3,)
+    assert len(running_stats._histograms) == 3
+    assert len(running_stats._bin_edges) == 3
+
+    # Verify basic statistics are reasonable
+    np.testing.assert_allclose(running_stats._mean, np.mean(data, axis=0), atol=1e-10)
+
+
+def test_running_quantile_stats_multiple_batch_updates():
+    """Test updating with multiple batches."""
+    np.random.seed(42)
+    data1 = np.random.normal(0, 1, (100, 2))
+    data2 = np.random.normal(1, 1, (150, 2))
+
+    running_stats = RunningQuantileStats()
+    running_stats.update(data1)
+    running_stats.update(data2)
+
+    assert running_stats._count == 250
+
+    # Verify running mean is correct
+    combined_data = np.vstack([data1, data2])
+    expected_mean = np.mean(combined_data, axis=0)
+    np.testing.assert_allclose(running_stats._mean, expected_mean, atol=1e-10)
+
+
+def test_running_quantile_stats_get_statistics_basic():
+    """Test getting basic statistics without quantiles."""
+    np.random.seed(42)
+    data = np.random.normal(0, 1, (100, 2))
+
+    running_stats = RunningQuantileStats()
+    running_stats.update(data)
+
+    stats = running_stats.get_statistics()
+
+    # Should have basic stats
+    expected_keys = {"min", "max", "mean", "std", "count"}
+    assert expected_keys.issubset(set(stats.keys()))
+
+    # Verify values
+    np.testing.assert_allclose(stats["mean"], np.mean(data, axis=0), atol=1e-10)
+    np.testing.assert_allclose(stats["std"], np.std(data, axis=0), atol=1e-6)
+    np.testing.assert_equal(stats["count"], np.array([100]))
+
+
+def test_running_quantile_stats_get_statistics_with_quantiles():
+    """Test getting statistics with quantiles."""
+    np.random.seed(42)
+    data = np.random.normal(0, 1, (1000, 2))
+
+    running_stats = RunningQuantileStats()
+    running_stats.update(data)
+
+    stats = running_stats.get_statistics()
+
+    # Should have basic stats plus quantiles
+    expected_keys = {"min", "max", "mean", "std", "count", "q01", "q10", "q50", "q90", "q99"}
+    assert expected_keys.issubset(set(stats.keys()))
+
+    # Verify quantile values are reasonable
+    from lerobot.datasets.compute_stats import DEFAULT_QUANTILES
+
+    for i, q in enumerate(DEFAULT_QUANTILES):
+        q_key = f"q{int(q * 100):02d}"
+        assert q_key in stats
+        assert stats[q_key].shape == (2,)
+
+        # Check that quantiles are in reasonable order
+        if i > 0:
+            prev_q_key = f"q{int(DEFAULT_QUANTILES[i - 1] * 100):02d}"
+            assert np.all(stats[prev_q_key] <= stats[q_key])
+
+
+def test_running_quantile_stats_histogram_adjustment():
+    """Test that histograms adjust when min/max change."""
+    running_stats = RunningQuantileStats()
+
+    # Initial data with small range
+    data1 = np.array([[0.0, 1.0], [0.1, 1.1], [0.2, 1.2]])
+    running_stats.update(data1)
+
+    initial_edges_0 = running_stats._bin_edges[0].copy()
+    initial_edges_1 = running_stats._bin_edges[1].copy()
+
+    # Add data with much larger range
+    data2 = np.array([[10.0, -10.0], [11.0, -11.0]])
+    running_stats.update(data2)
+
+    # Bin edges should have changed
+    assert not np.array_equal(initial_edges_0, running_stats._bin_edges[0])
+    assert not np.array_equal(initial_edges_1, running_stats._bin_edges[1])
+
+    # New edges should cover the expanded range
+    # First dimension: min should still be ~0.0, max should be ~11.0
+    assert running_stats._bin_edges[0][0] <= 0.0
+    assert running_stats._bin_edges[0][-1] >= 11.0
+
+    # Second dimension: min should be ~-11.0, max should be ~1.2
+    assert running_stats._bin_edges[1][0] <= -11.0
+    assert running_stats._bin_edges[1][-1] >= 1.2
+
+
+def test_running_quantile_stats_insufficient_data_error():
+    """Test error when trying to get stats with insufficient data."""
+    running_stats = RunningQuantileStats()
+
+    with pytest.raises(ValueError, match="Cannot compute statistics for less than 2 vectors"):
+        running_stats.get_statistics()
+
+    # Single vector should also fail
+    running_stats.update(np.array([[1.0]]))
+    with pytest.raises(ValueError, match="Cannot compute statistics for less than 2 vectors"):
+        running_stats.get_statistics()
+
+
+def test_running_quantile_stats_vector_length_consistency():
+    """Test error when vector lengths don't match."""
+    running_stats = RunningQuantileStats()
+    running_stats.update(np.array([[1.0, 2.0], [3.0, 4.0]]))
+
+    with pytest.raises(ValueError, match="The length of new vectors does not match"):
+        running_stats.update(np.array([[1.0, 2.0, 3.0]]))  # Different length
+
+
+def test_running_quantile_stats_reshape_handling():
+    """Test that various input shapes are handled correctly."""
+    running_stats = RunningQuantileStats()
+
+    # Test 3D input (e.g., images)
+    data_3d = np.random.normal(0, 1, (10, 32, 32))
+    running_stats.update(data_3d)
+
+    assert running_stats._count == 10 * 32
+    assert running_stats._mean.shape == (32,)
+
+    # Test 1D input
+    running_stats_1d = RunningQuantileStats()
+    data_1d = np.array([1, 2, 3, 4, 5]).reshape(-1, 1)
+    running_stats_1d.update(data_1d)
+
+    assert running_stats_1d._count == 5
+    assert running_stats_1d._mean.shape == (1,)
+
+
+def test_get_feature_stats_quantiles_enabled_by_default():
+    """Test that quantiles are computed by default."""
+    data = np.random.normal(0, 1, (100, 5))
+    stats = get_feature_stats(data, axis=0, keepdims=False)
+
+    expected_keys = {"min", "max", "mean", "std", "count", "q01", "q10", "q50", "q90", "q99"}
+    assert set(stats.keys()) == expected_keys
+
+
+def test_get_feature_stats_quantiles_with_vector_data():
+    """Test quantile computation with vector data."""
+    np.random.seed(42)
+    data = np.random.normal(0, 1, (100, 5))
+
+    stats = get_feature_stats(data, axis=0, keepdims=False)
+
+    expected_keys = {"min", "max", "mean", "std", "count", "q01", "q10", "q50", "q90", "q99"}
+    assert set(stats.keys()) == expected_keys
+
+    # Verify shapes
+    assert stats["q01"].shape == (5,)
+    assert stats["q99"].shape == (5,)
+
+    # Verify quantiles are reasonable
+    assert np.all(stats["q01"] < stats["q99"])
+
+
+def test_get_feature_stats_quantiles_with_image_data():
+    """Test quantile computation with image data."""
+    np.random.seed(42)
+    data = np.random.normal(0, 1, (50, 3, 32, 32))  # batch, channels, height, width
+
+    stats = get_feature_stats(data, axis=(0, 2, 3), keepdims=True)
+
+    expected_keys = {"min", "max", "mean", "std", "count", "q01", "q10", "q50", "q90", "q99"}
+    assert set(stats.keys()) == expected_keys
+
+    # Verify shapes for images (should be (1, channels, 1, 1))
+    assert stats["q01"].shape == (1, 3, 1, 1)
+    assert stats["q50"].shape == (1, 3, 1, 1)
+    assert stats["q99"].shape == (1, 3, 1, 1)
+
+
+def test_get_feature_stats_fixed_quantiles():
+    """Test that fixed quantiles are always computed."""
+    data = np.random.normal(0, 1, (200, 3))
+
+    stats = get_feature_stats(data, axis=0, keepdims=False)
+
+    expected_quantile_keys = {"q01", "q10", "q50", "q90", "q99"}
+    assert expected_quantile_keys.issubset(set(stats.keys()))
+
+
+def test_get_feature_stats_unsupported_axis_error():
+    """Test error for unsupported axis configuration."""
+    data = np.random.normal(0, 1, (10, 5))
+
+    with pytest.raises(ValueError, match="Unsupported axis configuration"):
+        get_feature_stats(
+            data,
+            axis=(1, 2),  # Unsupported axis
+            keepdims=False,
+        )
+
+
+def test_compute_episode_stats_backward_compatibility():
+    """Test that existing functionality is preserved."""
+    episode_data = {
+        "action": np.random.normal(0, 1, (100, 7)),
+        "observation.state": np.random.normal(0, 1, (100, 10)),
+    }
+    features = {
+        "action": {"dtype": "float32", "shape": (7,)},
+        "observation.state": {"dtype": "float32", "shape": (10,)},
+    }
+
+    stats = compute_episode_stats(episode_data, features)
+
+    for key in ["action", "observation.state"]:
+        expected_keys = {"min", "max", "mean", "std", "count", "q01", "q10", "q50", "q90", "q99"}
+        assert set(stats[key].keys()) == expected_keys
+
+
+def test_compute_episode_stats_with_custom_quantiles():
+    """Test quantile computation with custom quantile values."""
+    np.random.seed(42)
+    episode_data = {
+        "action": np.random.normal(0, 1, (100, 7)),
+        "observation.state": np.random.normal(2, 1, (100, 10)),
+    }
+    features = {
+        "action": {"dtype": "float32", "shape": (7,)},
+        "observation.state": {"dtype": "float32", "shape": (10,)},
+    }
+
+    stats = compute_episode_stats(episode_data, features)
+
+    # Should have quantiles
+    for key in ["action", "observation.state"]:
+        expected_keys = {"min", "max", "mean", "std", "count", "q01", "q10", "q50", "q90", "q99"}
+        assert set(stats[key].keys()) == expected_keys
+
+        # Verify shapes
+        assert stats[key]["q01"].shape == (features[key]["shape"][0],)
+        assert stats[key]["q99"].shape == (features[key]["shape"][0],)
+
+
+def test_compute_episode_stats_with_image_data():
+    """Test quantile computation with image features."""
+    image_paths = [f"image_{i}.jpg" for i in range(50)]
+    episode_data = {
+        "observation.image": image_paths,
+        "action": np.random.normal(0, 1, (50, 5)),
+    }
+    features = {
+        "observation.image": {"dtype": "image"},
+        "action": {"dtype": "float32", "shape": (5,)},
+    }
+
+    with patch("lerobot.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy):
+        stats = compute_episode_stats(episode_data, features)
+
+    # Image quantiles should be normalized and have correct shape
+    assert "q01" in stats["observation.image"]
+    assert "q50" in stats["observation.image"]
+    assert "q99" in stats["observation.image"]
+    assert stats["observation.image"]["q01"].shape == (3, 1, 1)
+    assert stats["observation.image"]["q50"].shape == (3, 1, 1)
+    assert stats["observation.image"]["q99"].shape == (3, 1, 1)
+
+    # Action quantiles should have correct shape
+    assert stats["action"]["q01"].shape == (5,)
+    assert stats["action"]["q50"].shape == (5,)
+    assert stats["action"]["q99"].shape == (5,)
+
+
+def test_compute_episode_stats_string_features_skipped():
+    """Test that string features are properly skipped."""
+    episode_data = {
+        "task": ["pick_apple"] * 100,  # String feature
+        "action": np.random.normal(0, 1, (100, 5)),
+    }
+    features = {
+        "task": {"dtype": "string"},
+        "action": {"dtype": "float32", "shape": (5,)},
+    }
+
+    stats = compute_episode_stats(
+        episode_data,
+        features,
+    )
+
+    # String features should be skipped
+    assert "task" not in stats
+    assert "action" in stats
+    assert "q01" in stats["action"]
+
+
+def test_aggregate_feature_stats_with_quantiles():
+    """Test aggregating feature stats that include quantiles."""
+    stats_ft_list = [
+        {
+            "min": np.array([1.0]),
+            "max": np.array([10.0]),
+            "mean": np.array([5.0]),
+            "std": np.array([2.0]),
+            "count": np.array([100]),
+            "q01": np.array([1.5]),
+            "q99": np.array([9.5]),
+        },
+        {
+            "min": np.array([2.0]),
+            "max": np.array([12.0]),
+            "mean": np.array([6.0]),
+            "std": np.array([2.5]),
+            "count": np.array([150]),
+            "q01": np.array([2.5]),
+            "q99": np.array([11.5]),
+        },
+    ]
+
+    result = aggregate_feature_stats(stats_ft_list)
+
+    # Should preserve quantiles
+    assert "q01" in result
+    assert "q99" in result
+
+    # Verify quantile aggregation (weighted average)
+    expected_q01 = (1.5 * 100 + 2.5 * 150) / 250  # ≈ 2.1
+    expected_q99 = (9.5 * 100 + 11.5 * 150) / 250  # ≈ 10.7
+
+    np.testing.assert_allclose(result["q01"], np.array([expected_q01]), atol=1e-6)
+    np.testing.assert_allclose(result["q99"], np.array([expected_q99]), atol=1e-6)
+
+
+def test_aggregate_stats_mixed_quantiles():
+    """Test aggregating stats where some have quantiles and some don't."""
+    stats_with_quantiles = {
+        "feature1": {
+            "min": np.array([1.0]),
+            "max": np.array([10.0]),
+            "mean": np.array([5.0]),
+            "std": np.array([2.0]),
+            "count": np.array([100]),
+            "q01": np.array([1.5]),
+            "q99": np.array([9.5]),
+        }
+    }
+
+    stats_without_quantiles = {
+        "feature2": {
+            "min": np.array([0.0]),
+            "max": np.array([5.0]),
+            "mean": np.array([2.5]),
+            "std": np.array([1.5]),
+            "count": np.array([50]),
+        }
+    }
+
+    all_stats = [stats_with_quantiles, stats_without_quantiles]
+    result = aggregate_stats(all_stats)
+
+    # Feature1 should keep its quantiles
+    assert "q01" in result["feature1"]
+    assert "q99" in result["feature1"]
+
+    # Feature2 should not have quantiles
+    assert "q01" not in result["feature2"]
+    assert "q99" not in result["feature2"]
+
+
+def test_assert_type_and_shape_with_quantiles():
+    """Test validation works correctly with quantile keys."""
+    # Valid stats with quantiles
+    valid_stats = [
+        {
+            "observation.image": {
+                "min": np.array([0.0, 0.0, 0.0]).reshape(3, 1, 1),
+                "max": np.array([1.0, 1.0, 1.0]).reshape(3, 1, 1),
+                "mean": np.array([0.5, 0.5, 0.5]).reshape(3, 1, 1),
+                "std": np.array([0.2, 0.2, 0.2]).reshape(3, 1, 1),
+                "count": np.array([100]),
+                "q01": np.array([0.1, 0.1, 0.1]).reshape(3, 1, 1),
+                "q99": np.array([0.9, 0.9, 0.9]).reshape(3, 1, 1),
+            }
+        }
+    ]
+
+    # Should not raise error
+    _assert_type_and_shape(valid_stats)
+
+    # Invalid shape for quantile
+    invalid_stats = [
+        {
+            "observation.image": {
+                "count": np.array([100]),
+                "q01": np.array([0.1, 0.2]),  # Wrong shape for image quantile
+            }
+        }
+    ]
+
+    with pytest.raises(ValueError, match="Shape of quantile 'q01' must be \\(3,1,1\\)"):
+        _assert_type_and_shape(invalid_stats)
+
+
+def test_quantile_integration_single_value_quantiles():
+    """Test quantile computation with single repeated value."""
+    data = np.ones((100, 3))  # All ones
+
+    running_stats = RunningQuantileStats()
+    running_stats.update(data)
+
+    stats = running_stats.get_statistics()
+
+    # All quantiles should be approximately 1.0
+    np.testing.assert_allclose(stats["q01"], np.array([1.0, 1.0, 1.0]), atol=1e-6)
+    np.testing.assert_allclose(stats["q50"], np.array([1.0, 1.0, 1.0]), atol=1e-6)
+    np.testing.assert_allclose(stats["q99"], np.array([1.0, 1.0, 1.0]), atol=1e-6)
+
+
+def test_quantile_integration_fixed_quantiles():
+    """Test that fixed quantiles are computed."""
+    np.random.seed(42)
+    data = np.random.normal(0, 1, (1000, 2))
+
+    stats = get_feature_stats(data, axis=0, keepdims=False)
+
+    # Check all fixed quantiles are present
+    assert "q01" in stats
+    assert "q10" in stats
+    assert "q50" in stats
+    assert "q90" in stats
+    assert "q99" in stats
+
+
+def test_quantile_integration_large_dataset_quantiles():
+    """Test quantile computation efficiency with large datasets."""
+    np.random.seed(42)
+    large_data = np.random.normal(0, 1, (10000, 5))
+
+    running_stats = RunningQuantileStats(num_quantile_bins=1000)  # Reduced bins for speed
+    running_stats.update(large_data)
+
+    stats = running_stats.get_statistics()
+
+    # Should complete without issues and produce reasonable results
+    assert stats["count"][0] == 10000
+    assert len(stats["q01"]) == 5
+
+
+def test_fixed_quantiles_always_computed():
+    """Test that the fixed quantiles [0.01, 0.10, 0.50, 0.90, 0.99] are always computed."""
+    np.random.seed(42)
+    # Test with vector data
+    vector_data = np.random.normal(0, 1, (100, 5))
+    vector_stats = get_feature_stats(vector_data, axis=0, keepdims=False)
+
+    # Check all fixed quantiles are present
+    expected_quantiles = ["q01", "q10", "q50", "q90", "q99"]
+    for q_key in expected_quantiles:
+        assert q_key in vector_stats
+        assert vector_stats[q_key].shape == (5,)
+
+    # Test with image data
+    image_data = np.random.randint(0, 256, (50, 3, 32, 32), dtype=np.uint8)
+    image_stats = get_feature_stats(image_data, axis=(0, 2, 3), keepdims=True)
+
+    # Check all fixed quantiles are present for images
+    for q_key in expected_quantiles:
+        assert q_key in image_stats
+        assert image_stats[q_key].shape == (1, 3, 1, 1)
+
+    # Test with episode data
+    episode_data = {
+        "action": np.random.normal(0, 1, (100, 7)),
+        "observation.state": np.random.normal(0, 1, (100, 10)),
+    }
+    features = {
+        "action": {"dtype": "float32", "shape": (7,)},
+        "observation.state": {"dtype": "float32", "shape": (10,)},
+    }
+
+    episode_stats = compute_episode_stats(episode_data, features)
+
+    # Check all fixed quantiles are present in episode stats
+    for key in ["action", "observation.state"]:
+        for q_key in expected_quantiles:
+            assert q_key in episode_stats[key]
+            assert episode_stats[key][q_key].shape == (features[key]["shape"][0],)
--- a/tests/datasets/test_quantiles_dataset_integration.py
+++ b/tests/datasets/test_quantiles_dataset_integration.py
@@ -0,0 +1,212 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Integration tests for quantile functionality in LeRobotDataset."""
+
+import numpy as np
+import pytest
+
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+
+
+def mock_load_image_as_numpy(path, dtype, channel_first):
+    """Mock image loading for consistent test results."""
+    return np.ones((3, 32, 32), dtype=dtype) if channel_first else np.ones((32, 32, 3), dtype=dtype)
+
+
+@pytest.fixture
+def simple_features():
+    """Simple feature configuration for testing."""
+    return {
+        "action": {
+            "dtype": "float32",
+            "shape": (4,),
+            "names": ["arm_x", "arm_y", "arm_z", "gripper"],
+        },
+        "observation.state": {
+            "dtype": "float32",
+            "shape": (10,),
+            "names": [f"joint_{i}" for i in range(10)],
+        },
+    }
+
+
+def test_create_dataset_with_fixed_quantiles(tmp_path, simple_features):
+    """Test creating dataset with fixed quantiles."""
+    dataset = LeRobotDataset.create(
+        repo_id="test_dataset_fixed_quantiles",
+        fps=30,
+        features=simple_features,
+        root=tmp_path / "create_fixed_quantiles",
+    )
+
+    # Dataset should be created successfully
+    assert dataset is not None
+
+
+def test_save_episode_computes_all_quantiles(tmp_path, simple_features):
+    """Test that all fixed quantiles are computed when saving an episode."""
+    dataset = LeRobotDataset.create(
+        repo_id="test_dataset_save_episode",
+        fps=30,
+        features=simple_features,
+        root=tmp_path / "save_episode_quantiles",
+    )
+
+    # Add some frames
+    for _ in range(10):
+        dataset.add_frame(
+            {
+                "action": np.random.randn(4).astype(np.float32),  # Correct shape for action
+                "observation.state": np.random.randn(10).astype(np.float32),
+                "task": "test_task",
+            }
+        )
+
+    dataset.save_episode()
+
+    # Check that all fixed quantiles were computed
+    stats = dataset.meta.stats
+    for key in ["action", "observation.state"]:
+        assert "q01" in stats[key]
+        assert "q10" in stats[key]
+        assert "q50" in stats[key]
+        assert "q90" in stats[key]
+        assert "q99" in stats[key]
+
+
+def test_quantile_values_ordering(tmp_path, simple_features):
+    """Test that quantile values are properly ordered."""
+    dataset = LeRobotDataset.create(
+        repo_id="test_dataset_quantile_ordering",
+        fps=30,
+        features=simple_features,
+        root=tmp_path / "quantile_ordering",
+    )
+
+    # Add data with known distribution
+    np.random.seed(42)
+    for _ in range(100):
+        dataset.add_frame(
+            {
+                "action": np.random.randn(4).astype(np.float32),  # Correct shape for action
+                "observation.state": np.random.randn(10).astype(np.float32),
+                "task": "test_task",
+            }
+        )
+
+    dataset.save_episode()
+    stats = dataset.meta.stats
+
+    # Verify quantile ordering
+    for key in ["action", "observation.state"]:
+        assert np.all(stats[key]["q01"] <= stats[key]["q10"])
+        assert np.all(stats[key]["q10"] <= stats[key]["q50"])
+        assert np.all(stats[key]["q50"] <= stats[key]["q90"])
+        assert np.all(stats[key]["q90"] <= stats[key]["q99"])
+
+
+def test_save_episode_with_fixed_quantiles(tmp_path, simple_features):
+    """Test saving episode always computes fixed quantiles."""
+    dataset = LeRobotDataset.create(
+        repo_id="test_dataset_save_fixed",
+        fps=30,
+        features=simple_features,
+        root=tmp_path / "save_fixed_quantiles",
+    )
+
+    # Add frames to episode
+    np.random.seed(42)
+    for _ in range(50):
+        frame = {
+            "action": np.random.normal(0, 1, (4,)).astype(np.float32),
+            "observation.state": np.random.normal(0, 1, (10,)).astype(np.float32),
+            "task": "test_task",
+        }
+        dataset.add_frame(frame)
+
+    dataset.save_episode()
+
+    # Check that all fixed quantiles are included
+    stats = dataset.meta.stats
+    for key in ["action", "observation.state"]:
+        feature_stats = stats[key]
+        expected_keys = {"min", "max", "mean", "std", "count", "q01", "q10", "q50", "q90", "q99"}
+        assert set(feature_stats.keys()) == expected_keys
+
+
+def test_quantile_aggregation_across_episodes(tmp_path, simple_features):
+    """Test quantile aggregation across multiple episodes."""
+    dataset = LeRobotDataset.create(
+        repo_id="test_dataset_aggregation",
+        fps=30,
+        features=simple_features,
+        root=tmp_path / "quantile_aggregation",
+    )
+
+    # Add frames to episode
+    np.random.seed(42)
+    for _ in range(100):
+        frame = {
+            "action": np.random.normal(0, 1, (4,)).astype(np.float32),
+            "observation.state": np.random.normal(2, 1, (10,)).astype(np.float32),
+            "task": "test_task",
+        }
+        dataset.add_frame(frame)
+
+    dataset.save_episode()
+
+    # Check stats include all fixed quantiles
+    stats = dataset.meta.stats
+    for key in ["action", "observation.state"]:
+        feature_stats = stats[key]
+        expected_keys = {"min", "max", "mean", "std", "count", "q01", "q10", "q50", "q90", "q99"}
+        assert set(feature_stats.keys()) == expected_keys
+        assert feature_stats["q01"].shape == (simple_features[key]["shape"][0],)
+        assert feature_stats["q50"].shape == (simple_features[key]["shape"][0],)
+        assert feature_stats["q99"].shape == (simple_features[key]["shape"][0],)
+        assert np.all(feature_stats["q01"] <= feature_stats["q50"])
+        assert np.all(feature_stats["q50"] <= feature_stats["q99"])
+
+
+def test_save_multiple_episodes_with_quantiles(tmp_path, simple_features):
+    """Test quantile aggregation across multiple episodes."""
+    dataset = LeRobotDataset.create(
+        repo_id="test_dataset_multiple_episodes",
+        fps=30,
+        features=simple_features,
+        root=tmp_path / "multiple_episodes",
+    )
+
+    # Save multiple episodes
+    np.random.seed(42)
+    for episode_idx in range(3):
+        for _ in range(50):
+            frame = {
+                "action": np.random.normal(episode_idx * 2.0, 1, (4,)).astype(np.float32),
+                "observation.state": np.random.normal(-episode_idx * 1.5, 1, (10,)).astype(np.float32),
+                "task": f"task_{episode_idx}",
+            }
+            dataset.add_frame(frame)
+
+        dataset.save_episode()
+
+    # Verify final stats include properly aggregated quantiles
+    stats = dataset.meta.stats
+    for key in ["action", "observation.state"]:
+        feature_stats = stats[key]
+        assert "q01" in feature_stats and "q99" in feature_stats
+        assert feature_stats["count"][0] == 150  # 3 episodes * 50 frames