Incorrect weighted aggregation of q01/q99 when merging dataset statistics

[asmashine]

Ticket Type

🐛 Bug Report (Something isn't working)

Environment & System Info

- LeRobot version: 0.6.0
- Platform: Linux-6.8.0-124-generic-x86_64-with-glibc2.35
- Python version: 3.12.13
- Huggingface Hub version: 1.23.0
- Transformers version: 5.5.4
- Datasets version: 4.8.5
- Numpy version: 2.2.6
- FFmpeg version: N/A
- PyTorch version: 2.11.0+cu128
- Torchcodec version: 0.11.1+cpu
- Is PyTorch built with CUDA support?: True
- Cuda version: 12.8
- GPU model: NVIDIA GeForce RTX 5090
- Using GPU in script?: <fill in>
- lerobot scripts: ['lerobot-annotate', 'lerobot-calibrate', 'lerobot-dataset-viz', 'lerobot-edit-dataset', 'lerobot-eval', 'lerobot-find-cameras', 'lerobot-find-joint-limits', 'lerobot-find-port', 'lerobot-imgtransform-viz', 'lerobot-info', 'lerobot-record', 'lerobot-replay', 'lerobot-rollout', 'lerobot-setup-can', 'lerobot-setup-motors', 'lerobot-teleoperate', 'lerobot-train', 'lerobot-train-tokenizer']

Description

Incorrect quantile aggregation when merging dataset statistics

When statistics from multiple datasets or data sources are merged, LeRobot correctly aggregates mean and std, but incorrectly computes quantiles such as q01 and q99 by taking a frame-count-weighted average of the per-source quantiles.

Problematic code

File:

src/lerobot/datasets/compute_stats.py

Problematic logic in aggregate_feature_stats:

for q_key in quantile_keys:
    if all(q_key in s for s in stats_ft_list):
        quantile_values = np.stack([s[q_key] for s in stats_ft_list])
        weighted_quantiles = quantile_values * counts
        aggregated[q_key] = weighted_quantiles.sum(axis=0) / total_count

This effectively computes:

merged_q01 = Σ(weight_i × source_q01_i)
merged_q99 = Σ(weight_i × source_q99_i)

However, quantiles are not linearly composable. A weighted average of per-source quantiles is generally not equal to the quantile of the combined samples.

For datasets whose sources cover different regions of the state or action space, this can produce severely incorrect and often overly narrow q01/q99 ranges.

In our merged dataset, recomputing the quantiles directly over all training frames produced the following results:

Feature              Outside saved q01/q99
state joint 1        41.65%
state joint 6        38.74%
action joint 1       42.12%
action joint 6       38.83%

The differences in mean were only around 1e-7 to 1e-8, confirming that the issue is specific to quantile aggregation rather than data loading or mean/std aggregation.

This significantly affects policies using QUANTILES normalization, such as Pi0.5, because a large fraction of normal training samples are incorrectly mapped outside the expected [-1, 1] range.

Expected behavior

Quantiles should be computed directly from all samples in the final training dataset, for example:

q01, q99 = np.quantile(
    all_train_values,
    [0.01, 0.99],
    axis=0,
    method="linear",
)

Alternatively, aggregation should use a mergeable quantile estimator such as a quantile sketch. If neither the original samples nor a mergeable estimator is available, the code should not silently generate quantiles by averaging per-source values.

Context & Reproduction

No response

Relevant logs or stack trace

Checklist

• I have searched existing tickets to ensure this isn't a duplicate.
• I am using the latest version of the main branch.
• I have verified this is not an environment-specific problem.

Additional Info / Workarounds

No response