mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-14 16:19:45 +00:00
Steven Palma
837 lines
28 KiB
Python
837 lines
28 KiB
Python
#!/usr/bin/env python
|
|
|
|
# Copyright 2024 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.
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
|
|
|
from lerobot.datasets.compute_stats import (
|
|
RunningQuantileStats,
|
|
_assert_type_and_shape,
|
|
aggregate_feature_stats,
|
|
aggregate_stats,
|
|
compute_episode_stats,
|
|
estimate_num_samples,
|
|
get_feature_stats,
|
|
sample_images,
|
|
sample_indices,
|
|
)
|
|
from lerobot.utils.constants import OBS_IMAGE, OBS_STATE
|
|
|
|
|
|
def mock_load_image_as_numpy(path, dtype, channel_first):
|
|
return np.ones((3, 32, 32), dtype=dtype) if channel_first else np.ones((32, 32, 3), dtype=dtype)
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_array():
|
|
return np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
|
|
|
|
|
|
def test_estimate_num_samples():
|
|
assert estimate_num_samples(1) == 1
|
|
assert estimate_num_samples(10) == 10
|
|
assert estimate_num_samples(100) == 100
|
|
assert estimate_num_samples(200) == 100
|
|
assert estimate_num_samples(1000) == 177
|
|
assert estimate_num_samples(2000) == 299
|
|
assert estimate_num_samples(5000) == 594
|
|
assert estimate_num_samples(10_000) == 1000
|
|
assert estimate_num_samples(20_000) == 1681
|
|
assert estimate_num_samples(50_000) == 3343
|
|
assert estimate_num_samples(500_000) == 10_000
|
|
|
|
|
|
def test_sample_indices():
|
|
indices = sample_indices(10)
|
|
assert len(indices) > 0
|
|
assert indices[0] == 0
|
|
assert indices[-1] == 9
|
|
assert len(indices) == estimate_num_samples(10)
|
|
|
|
|
|
@patch("lerobot.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy)
|
|
def test_sample_images(mock_load):
|
|
image_paths = [f"image_{i}.jpg" for i in range(100)]
|
|
images = sample_images(image_paths)
|
|
assert isinstance(images, np.ndarray)
|
|
assert images.shape[1:] == (3, 32, 32)
|
|
assert images.dtype == np.uint8
|
|
assert len(images) == estimate_num_samples(100)
|
|
|
|
|
|
def test_get_feature_stats_images():
|
|
data = np.random.rand(100, 3, 32, 32)
|
|
stats = get_feature_stats(data, axis=(0, 2, 3), keepdims=True)
|
|
assert "min" in stats and "max" in stats and "mean" in stats and "std" in stats and "count" in stats
|
|
np.testing.assert_equal(stats["count"], np.array([100]))
|
|
assert stats["min"].shape == stats["max"].shape == stats["mean"].shape == stats["std"].shape
|
|
|
|
|
|
def test_get_feature_stats_axis_0_keepdims(sample_array):
|
|
expected = {
|
|
"min": np.array([[1, 2, 3]]),
|
|
"max": np.array([[7, 8, 9]]),
|
|
"mean": np.array([[4.0, 5.0, 6.0]]),
|
|
"std": np.array([[2.44948974, 2.44948974, 2.44948974]]),
|
|
"count": np.array([3]),
|
|
}
|
|
result = get_feature_stats(sample_array, axis=(0,), keepdims=True)
|
|
for key in expected:
|
|
np.testing.assert_allclose(result[key], expected[key])
|
|
|
|
|
|
def test_get_feature_stats_axis_1(sample_array):
|
|
expected = {
|
|
"min": np.array([1, 4, 7]),
|
|
"max": np.array([3, 6, 9]),
|
|
"mean": np.array([2.0, 5.0, 8.0]),
|
|
"std": np.array([0.81649658, 0.81649658, 0.81649658]),
|
|
"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])
|
|
|
|
|
|
def test_get_feature_stats_no_axis(sample_array):
|
|
expected = {
|
|
"min": np.array(1),
|
|
"max": np.array(9),
|
|
"mean": np.array(5.0),
|
|
"std": np.array(2.5819889),
|
|
"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])
|
|
|
|
|
|
def test_get_feature_stats_empty_array():
|
|
array = np.array([])
|
|
with pytest.raises(ValueError):
|
|
get_feature_stats(array, axis=(0,), keepdims=True)
|
|
|
|
|
|
def test_get_feature_stats_single_value():
|
|
array = np.array([[1337]])
|
|
result = get_feature_stats(array, axis=None, keepdims=True)
|
|
np.testing.assert_equal(result["min"], np.array(1337))
|
|
np.testing.assert_equal(result["max"], np.array(1337))
|
|
np.testing.assert_equal(result["mean"], np.array(1337.0))
|
|
np.testing.assert_equal(result["std"], np.array(0.0))
|
|
np.testing.assert_equal(result["count"], np.array([1]))
|
|
|
|
|
|
def test_compute_episode_stats():
|
|
episode_data = {
|
|
OBS_IMAGE: [f"image_{i}.jpg" for i in range(100)],
|
|
OBS_STATE: np.random.rand(100, 10),
|
|
}
|
|
features = {
|
|
OBS_IMAGE: {"dtype": "image"},
|
|
OBS_STATE: {"dtype": "numeric"},
|
|
}
|
|
|
|
with patch("lerobot.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy):
|
|
stats = compute_episode_stats(episode_data, features)
|
|
|
|
assert OBS_IMAGE in stats and OBS_STATE in stats
|
|
assert stats[OBS_IMAGE]["count"].item() == 100
|
|
assert stats[OBS_STATE]["count"].item() == 100
|
|
assert stats[OBS_IMAGE]["mean"].shape == (3, 1, 1)
|
|
|
|
|
|
def test_assert_type_and_shape_valid():
|
|
valid_stats = [
|
|
{
|
|
"feature1": {
|
|
"min": np.array([1.0]),
|
|
"max": np.array([10.0]),
|
|
"mean": np.array([5.0]),
|
|
"std": np.array([2.0]),
|
|
"count": np.array([1]),
|
|
}
|
|
}
|
|
]
|
|
_assert_type_and_shape(valid_stats)
|
|
|
|
|
|
def test_assert_type_and_shape_invalid_type():
|
|
invalid_stats = [
|
|
{
|
|
"feature1": {
|
|
"min": [1.0], # Not a numpy array
|
|
"max": np.array([10.0]),
|
|
"mean": np.array([5.0]),
|
|
"std": np.array([2.0]),
|
|
"count": np.array([1]),
|
|
}
|
|
}
|
|
]
|
|
with pytest.raises(ValueError, match="Stats must be composed of numpy array"):
|
|
_assert_type_and_shape(invalid_stats)
|
|
|
|
|
|
def test_assert_type_and_shape_invalid_shape():
|
|
invalid_stats = [
|
|
{
|
|
"feature1": {
|
|
"count": np.array([1, 2]), # Wrong shape
|
|
}
|
|
}
|
|
]
|
|
with pytest.raises(ValueError, match=r"Shape of 'count' must be \(1\)"):
|
|
_assert_type_and_shape(invalid_stats)
|
|
|
|
|
|
def test_aggregate_feature_stats():
|
|
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([1]),
|
|
},
|
|
{
|
|
"min": np.array([2.0]),
|
|
"max": np.array([12.0]),
|
|
"mean": np.array([6.0]),
|
|
"std": np.array([2.5]),
|
|
"count": np.array([1]),
|
|
},
|
|
]
|
|
result = aggregate_feature_stats(stats_ft_list)
|
|
np.testing.assert_allclose(result["min"], np.array([1.0]))
|
|
np.testing.assert_allclose(result["max"], np.array([12.0]))
|
|
np.testing.assert_allclose(result["mean"], np.array([5.5]))
|
|
np.testing.assert_allclose(result["std"], np.array([2.318405]), atol=1e-6)
|
|
np.testing.assert_allclose(result["count"], np.array([2]))
|
|
|
|
|
|
def test_aggregate_stats():
|
|
all_stats = [
|
|
{
|
|
OBS_IMAGE: {
|
|
"min": [1, 2, 3],
|
|
"max": [10, 20, 30],
|
|
"mean": [5.5, 10.5, 15.5],
|
|
"std": [2.87, 5.87, 8.87],
|
|
"count": 10,
|
|
},
|
|
OBS_STATE: {"min": 1, "max": 10, "mean": 5.5, "std": 2.87, "count": 10},
|
|
"extra_key_0": {"min": 5, "max": 25, "mean": 15, "std": 6, "count": 6},
|
|
},
|
|
{
|
|
OBS_IMAGE: {
|
|
"min": [2, 1, 0],
|
|
"max": [15, 10, 5],
|
|
"mean": [8.5, 5.5, 2.5],
|
|
"std": [3.42, 2.42, 1.42],
|
|
"count": 15,
|
|
},
|
|
OBS_STATE: {"min": 2, "max": 15, "mean": 8.5, "std": 3.42, "count": 15},
|
|
"extra_key_1": {"min": 0, "max": 20, "mean": 10, "std": 5, "count": 5},
|
|
},
|
|
]
|
|
|
|
expected_agg_stats = {
|
|
OBS_IMAGE: {
|
|
"min": [1, 1, 0],
|
|
"max": [15, 20, 30],
|
|
"mean": [7.3, 7.5, 7.7],
|
|
"std": [3.5317, 4.8267, 8.5581],
|
|
"count": 25,
|
|
},
|
|
OBS_STATE: {
|
|
"min": 1,
|
|
"max": 15,
|
|
"mean": 7.3,
|
|
"std": 3.5317,
|
|
"count": 25,
|
|
},
|
|
"extra_key_0": {
|
|
"min": 5,
|
|
"max": 25,
|
|
"mean": 15.0,
|
|
"std": 6.0,
|
|
"count": 6,
|
|
},
|
|
"extra_key_1": {
|
|
"min": 0,
|
|
"max": 20,
|
|
"mean": 10.0,
|
|
"std": 5.0,
|
|
"count": 5,
|
|
},
|
|
}
|
|
|
|
# cast to numpy
|
|
for ep_stats in all_stats:
|
|
for fkey, stats in ep_stats.items():
|
|
for k in stats:
|
|
stats[k] = np.array(stats[k], dtype=np.int64 if k == "count" else np.float32)
|
|
if fkey == OBS_IMAGE and k != "count":
|
|
stats[k] = stats[k].reshape(3, 1, 1) # for normalization on image channels
|
|
else:
|
|
stats[k] = stats[k].reshape(1)
|
|
|
|
# cast to numpy
|
|
for fkey, stats in expected_agg_stats.items():
|
|
for k in stats:
|
|
stats[k] = np.array(stats[k], dtype=np.int64 if k == "count" else np.float32)
|
|
if fkey == OBS_IMAGE and k != "count":
|
|
stats[k] = stats[k].reshape(3, 1, 1) # for normalization on image channels
|
|
else:
|
|
stats[k] = stats[k].reshape(1)
|
|
|
|
results = aggregate_stats(all_stats)
|
|
|
|
for fkey in expected_agg_stats:
|
|
np.testing.assert_allclose(results[fkey]["min"], expected_agg_stats[fkey]["min"])
|
|
np.testing.assert_allclose(results[fkey]["max"], expected_agg_stats[fkey]["max"])
|
|
np.testing.assert_allclose(results[fkey]["mean"], expected_agg_stats[fkey]["mean"])
|
|
np.testing.assert_allclose(
|
|
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],)
|