257 lines
8.7 KiB
Python
257 lines
8.7 KiB
Python
import pytest
|
|
import torch
|
|
from datasets import Dataset
|
|
|
|
from lerobot.common.datasets.utils import (
|
|
calculate_episode_data_index,
|
|
check_delta_timestamps,
|
|
check_timestamps_sync,
|
|
get_delta_indices,
|
|
hf_transform_to_torch,
|
|
)
|
|
from tests.fixtures.constants import DUMMY_MOTOR_FEATURES
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def synced_hf_dataset_factory(hf_dataset_factory):
|
|
def _create_synced_hf_dataset(fps: int = 30) -> Dataset:
|
|
return hf_dataset_factory(fps=fps)
|
|
|
|
return _create_synced_hf_dataset
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def unsynced_hf_dataset_factory(synced_hf_dataset_factory):
|
|
def _create_unsynced_hf_dataset(fps: int = 30, tolerance_s: float = 1e-4) -> Dataset:
|
|
hf_dataset = synced_hf_dataset_factory(fps=fps)
|
|
features = hf_dataset.features
|
|
df = hf_dataset.to_pandas()
|
|
dtype = df["timestamp"].dtype # This is to avoid pandas type warning
|
|
# Modify a single timestamp just outside tolerance
|
|
df.at[30, "timestamp"] = dtype.type(df.at[30, "timestamp"] + (tolerance_s * 1.1))
|
|
unsynced_hf_dataset = Dataset.from_pandas(df, features=features)
|
|
unsynced_hf_dataset.set_transform(hf_transform_to_torch)
|
|
return unsynced_hf_dataset
|
|
|
|
return _create_unsynced_hf_dataset
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def slightly_off_hf_dataset_factory(synced_hf_dataset_factory):
|
|
def _create_slightly_off_hf_dataset(fps: int = 30, tolerance_s: float = 1e-4) -> Dataset:
|
|
hf_dataset = synced_hf_dataset_factory(fps=fps)
|
|
features = hf_dataset.features
|
|
df = hf_dataset.to_pandas()
|
|
dtype = df["timestamp"].dtype # This is to avoid pandas type warning
|
|
# Modify a single timestamp just inside tolerance
|
|
df.at[30, "timestamp"] = dtype.type(df.at[30, "timestamp"] + (tolerance_s * 0.9))
|
|
unsynced_hf_dataset = Dataset.from_pandas(df, features=features)
|
|
unsynced_hf_dataset.set_transform(hf_transform_to_torch)
|
|
return unsynced_hf_dataset
|
|
|
|
return _create_slightly_off_hf_dataset
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def valid_delta_timestamps_factory():
|
|
def _create_valid_delta_timestamps(fps: int = 30, keys: list = DUMMY_MOTOR_FEATURES) -> dict:
|
|
delta_timestamps = {key: [i * (1 / fps) for i in range(-10, 10)] for key in keys}
|
|
return delta_timestamps
|
|
|
|
return _create_valid_delta_timestamps
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def invalid_delta_timestamps_factory(valid_delta_timestamps_factory):
|
|
def _create_invalid_delta_timestamps(
|
|
fps: int = 30, tolerance_s: float = 1e-4, keys: list = DUMMY_MOTOR_FEATURES
|
|
) -> dict:
|
|
delta_timestamps = valid_delta_timestamps_factory(fps, keys)
|
|
# Modify a single timestamp just outside tolerance
|
|
for key in keys:
|
|
delta_timestamps[key][3] += tolerance_s * 1.1
|
|
return delta_timestamps
|
|
|
|
return _create_invalid_delta_timestamps
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def slightly_off_delta_timestamps_factory(valid_delta_timestamps_factory):
|
|
def _create_slightly_off_delta_timestamps(
|
|
fps: int = 30, tolerance_s: float = 1e-4, keys: list = DUMMY_MOTOR_FEATURES
|
|
) -> dict:
|
|
delta_timestamps = valid_delta_timestamps_factory(fps, keys)
|
|
# Modify a single timestamp just inside tolerance
|
|
for key in delta_timestamps:
|
|
delta_timestamps[key][3] += tolerance_s * 0.9
|
|
delta_timestamps[key][-3] += tolerance_s * 0.9
|
|
return delta_timestamps
|
|
|
|
return _create_slightly_off_delta_timestamps
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def delta_indices(keys: list = DUMMY_MOTOR_FEATURES) -> dict:
|
|
return {key: list(range(-10, 10)) for key in keys}
|
|
|
|
|
|
def test_check_timestamps_sync_synced(synced_hf_dataset_factory):
|
|
fps = 30
|
|
tolerance_s = 1e-4
|
|
synced_hf_dataset = synced_hf_dataset_factory(fps)
|
|
episode_data_index = calculate_episode_data_index(synced_hf_dataset)
|
|
result = check_timestamps_sync(
|
|
hf_dataset=synced_hf_dataset,
|
|
episode_data_index=episode_data_index,
|
|
fps=fps,
|
|
tolerance_s=tolerance_s,
|
|
)
|
|
assert result is True
|
|
|
|
|
|
def test_check_timestamps_sync_unsynced(unsynced_hf_dataset_factory):
|
|
fps = 30
|
|
tolerance_s = 1e-4
|
|
unsynced_hf_dataset = unsynced_hf_dataset_factory(fps, tolerance_s)
|
|
episode_data_index = calculate_episode_data_index(unsynced_hf_dataset)
|
|
with pytest.raises(ValueError):
|
|
check_timestamps_sync(
|
|
hf_dataset=unsynced_hf_dataset,
|
|
episode_data_index=episode_data_index,
|
|
fps=fps,
|
|
tolerance_s=tolerance_s,
|
|
)
|
|
|
|
|
|
def test_check_timestamps_sync_unsynced_no_exception(unsynced_hf_dataset_factory):
|
|
fps = 30
|
|
tolerance_s = 1e-4
|
|
unsynced_hf_dataset = unsynced_hf_dataset_factory(fps, tolerance_s)
|
|
episode_data_index = calculate_episode_data_index(unsynced_hf_dataset)
|
|
result = check_timestamps_sync(
|
|
hf_dataset=unsynced_hf_dataset,
|
|
episode_data_index=episode_data_index,
|
|
fps=fps,
|
|
tolerance_s=tolerance_s,
|
|
raise_value_error=False,
|
|
)
|
|
assert result is False
|
|
|
|
|
|
def test_check_timestamps_sync_slightly_off(slightly_off_hf_dataset_factory):
|
|
fps = 30
|
|
tolerance_s = 1e-4
|
|
slightly_off_hf_dataset = slightly_off_hf_dataset_factory(fps, tolerance_s)
|
|
episode_data_index = calculate_episode_data_index(slightly_off_hf_dataset)
|
|
result = check_timestamps_sync(
|
|
hf_dataset=slightly_off_hf_dataset,
|
|
episode_data_index=episode_data_index,
|
|
fps=fps,
|
|
tolerance_s=tolerance_s,
|
|
)
|
|
assert result is True
|
|
|
|
|
|
def test_check_timestamps_sync_single_timestamp():
|
|
single_timestamp_hf_dataset = Dataset.from_dict({"timestamp": [0.0], "episode_index": [0]})
|
|
single_timestamp_hf_dataset.set_transform(hf_transform_to_torch)
|
|
episode_data_index = {"to": torch.tensor([1]), "from": torch.tensor([0])}
|
|
fps = 30
|
|
tolerance_s = 1e-4
|
|
result = check_timestamps_sync(
|
|
hf_dataset=single_timestamp_hf_dataset,
|
|
episode_data_index=episode_data_index,
|
|
fps=fps,
|
|
tolerance_s=tolerance_s,
|
|
)
|
|
assert result is True
|
|
|
|
|
|
# TODO(aliberts): Change behavior of hf_transform_to_torch so that it can work with empty dataset
|
|
@pytest.mark.skip("TODO: fix")
|
|
def test_check_timestamps_sync_empty_dataset():
|
|
fps = 30
|
|
tolerance_s = 1e-4
|
|
empty_hf_dataset = Dataset.from_dict({"timestamp": [], "episode_index": []})
|
|
empty_hf_dataset.set_transform(hf_transform_to_torch)
|
|
episode_data_index = {
|
|
"to": torch.tensor([], dtype=torch.int64),
|
|
"from": torch.tensor([], dtype=torch.int64),
|
|
}
|
|
result = check_timestamps_sync(
|
|
hf_dataset=empty_hf_dataset,
|
|
episode_data_index=episode_data_index,
|
|
fps=fps,
|
|
tolerance_s=tolerance_s,
|
|
)
|
|
assert result is True
|
|
|
|
|
|
def test_check_delta_timestamps_valid(valid_delta_timestamps_factory):
|
|
fps = 30
|
|
tolerance_s = 1e-4
|
|
valid_delta_timestamps = valid_delta_timestamps_factory(fps)
|
|
result = check_delta_timestamps(
|
|
delta_timestamps=valid_delta_timestamps,
|
|
fps=fps,
|
|
tolerance_s=tolerance_s,
|
|
)
|
|
assert result is True
|
|
|
|
|
|
def test_check_delta_timestamps_slightly_off(slightly_off_delta_timestamps_factory):
|
|
fps = 30
|
|
tolerance_s = 1e-4
|
|
slightly_off_delta_timestamps = slightly_off_delta_timestamps_factory(fps, tolerance_s)
|
|
result = check_delta_timestamps(
|
|
delta_timestamps=slightly_off_delta_timestamps,
|
|
fps=fps,
|
|
tolerance_s=tolerance_s,
|
|
)
|
|
assert result is True
|
|
|
|
|
|
def test_check_delta_timestamps_invalid(invalid_delta_timestamps_factory):
|
|
fps = 30
|
|
tolerance_s = 1e-4
|
|
invalid_delta_timestamps = invalid_delta_timestamps_factory(fps, tolerance_s)
|
|
with pytest.raises(ValueError):
|
|
check_delta_timestamps(
|
|
delta_timestamps=invalid_delta_timestamps,
|
|
fps=fps,
|
|
tolerance_s=tolerance_s,
|
|
)
|
|
|
|
|
|
def test_check_delta_timestamps_invalid_no_exception(invalid_delta_timestamps_factory):
|
|
fps = 30
|
|
tolerance_s = 1e-4
|
|
invalid_delta_timestamps = invalid_delta_timestamps_factory(fps, tolerance_s)
|
|
result = check_delta_timestamps(
|
|
delta_timestamps=invalid_delta_timestamps,
|
|
fps=fps,
|
|
tolerance_s=tolerance_s,
|
|
raise_value_error=False,
|
|
)
|
|
assert result is False
|
|
|
|
|
|
def test_check_delta_timestamps_empty():
|
|
delta_timestamps = {}
|
|
fps = 30
|
|
tolerance_s = 1e-4
|
|
result = check_delta_timestamps(
|
|
delta_timestamps=delta_timestamps,
|
|
fps=fps,
|
|
tolerance_s=tolerance_s,
|
|
)
|
|
assert result is True
|
|
|
|
|
|
def test_delta_indices(valid_delta_timestamps_factory, delta_indices):
|
|
fps = 30
|
|
delta_timestamps = valid_delta_timestamps_factory(fps)
|
|
expected_delta_indices = delta_indices
|
|
actual_delta_indices = get_delta_indices(delta_timestamps, fps)
|
|
assert expected_delta_indices == actual_delta_indices
|