mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-07 01:51:47 +00:00
Merge branch 'main' into feat/add-recap
This commit is contained in:
Khalil Meftah
@@ -47,6 +47,7 @@ class _FakeMeta:
|
||||
def __init__(self, video_keys: list[str], image_keys: list[str], video_path: Path | None = None) -> None:
|
||||
self.video_keys = video_keys
|
||||
self.camera_keys = [*video_keys, *image_keys]
|
||||
self.depth_keys = []
|
||||
self._video_path = video_path
|
||||
self.episodes = {0: {f"videos/{key}/from_timestamp": 0.0 for key in video_keys}}
|
||||
|
||||
@@ -208,14 +209,14 @@ def test_episode_clip_path_trims_via_reencode_video(tmp_path: Path, monkeypatch)
|
||||
def fake_reencode(
|
||||
input_video_path,
|
||||
output_video_path,
|
||||
camera_encoder=None,
|
||||
video_encoder=None,
|
||||
overwrite=False,
|
||||
start_time_s=None,
|
||||
end_time_s=None,
|
||||
):
|
||||
captured.update(
|
||||
src=Path(input_video_path),
|
||||
encoder=camera_encoder,
|
||||
encoder=video_encoder,
|
||||
start_time_s=start_time_s,
|
||||
end_time_s=end_time_s,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
# 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.
|
||||
|
||||
import pytest
|
||||
|
||||
import lerobot.configs.train as tc
|
||||
from lerobot.configs.train import TrainPipelineConfig
|
||||
|
||||
|
||||
class _FakeHTTPError(tc.HfHubHTTPError):
|
||||
"""HfHubHTTPError that can be raised without a real HTTP response object."""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
|
||||
def test_from_pretrained_falls_back_to_latest_checkpoint_config(tmp_path, monkeypatch):
|
||||
"""A Hub repo with no root train_config.json (an interrupted run that only pushed
|
||||
checkpoints/) resolves via the latest checkpoint's config."""
|
||||
# A real train_config.json written by save_pretrained, to be returned by the fallback.
|
||||
parsed = tc.draccus.parse(TrainPipelineConfig, args=["--dataset.repo_id", "u/d"])
|
||||
cfg_file = tmp_path / "train_config.json"
|
||||
parsed._save_pretrained(tmp_path)
|
||||
assert cfg_file.is_file()
|
||||
|
||||
calls = []
|
||||
|
||||
def fake_hf_hub_download(filename=None, **kwargs):
|
||||
calls.append(filename)
|
||||
if filename == "train_config.json":
|
||||
raise _FakeHTTPError() # no root config
|
||||
if filename == "checkpoints/000010/pretrained_model/train_config.json":
|
||||
return str(cfg_file)
|
||||
raise AssertionError(f"unexpected filename {filename}")
|
||||
|
||||
monkeypatch.setattr(tc, "hf_hub_download", fake_hf_hub_download)
|
||||
monkeypatch.setattr(
|
||||
tc, "find_latest_hub_checkpoint", lambda repo_id, token=None, revision=None: "checkpoints/000010"
|
||||
)
|
||||
|
||||
loaded = TrainPipelineConfig.from_pretrained("user/interrupted-run")
|
||||
assert loaded.dataset.repo_id == "u/d"
|
||||
# Tried the root config first, then fell back to the latest checkpoint's config.
|
||||
assert calls == ["train_config.json", "checkpoints/000010/pretrained_model/train_config.json"]
|
||||
|
||||
|
||||
def test_from_pretrained_raises_when_no_root_config_and_no_checkpoints(monkeypatch):
|
||||
"""No root config AND no checkpoints → a clear FileNotFoundError, not the raw HTTP error."""
|
||||
|
||||
def fake_hf_hub_download(filename=None, **kwargs):
|
||||
raise _FakeHTTPError()
|
||||
|
||||
monkeypatch.setattr(tc, "hf_hub_download", fake_hf_hub_download)
|
||||
monkeypatch.setattr(tc, "find_latest_hub_checkpoint", lambda repo_id, token=None, revision=None: None)
|
||||
|
||||
with pytest.raises(FileNotFoundError, match="train_config.json not found"):
|
||||
TrainPipelineConfig.from_pretrained("user/empty-repo")
|
||||
@@ -29,7 +29,10 @@ from lerobot.configs import VIDEO_ENCODER_INFO_KEYS
|
||||
from lerobot.datasets.aggregate import aggregate_datasets
|
||||
from lerobot.datasets.feature_utils import features_equal_for_merge
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
from tests.fixtures.constants import DUMMY_REPO_ID
|
||||
from tests.fixtures.constants import (
|
||||
DUMMY_CAMERA_FEATURES_WITH_DEPTH,
|
||||
DUMMY_REPO_ID,
|
||||
)
|
||||
|
||||
|
||||
def assert_data_shards_one_row_group_per_episode(root):
|
||||
@@ -211,6 +214,26 @@ def assert_dataset_iteration_works(aggr_ds):
|
||||
pass
|
||||
|
||||
|
||||
def assert_depth_keys_preserved(aggr_ds, ds_0, ds_1):
|
||||
"""Test that depth keys are correctly preserved after aggregation.
|
||||
|
||||
Ensures that the ``is_depth_map`` marker on visual features survives
|
||||
aggregation, so that downstream consumers (e.g. the dataset reader's
|
||||
depth decoding path) keep working on the merged dataset.
|
||||
"""
|
||||
expected_depth_keys = set(ds_0.meta.depth_keys)
|
||||
assert expected_depth_keys == set(ds_1.meta.depth_keys), (
|
||||
"Source datasets disagree on depth_keys; test setup is inconsistent"
|
||||
)
|
||||
actual_depth_keys = set(aggr_ds.meta.depth_keys)
|
||||
assert actual_depth_keys == expected_depth_keys, (
|
||||
f"Expected depth_keys {expected_depth_keys}, got {actual_depth_keys}"
|
||||
)
|
||||
for key in expected_depth_keys:
|
||||
info = aggr_ds.meta.info.features[key].get("info") or {}
|
||||
assert info.get("is_depth_map") is True, f"Depth marker lost on feature {key!r} after aggregation"
|
||||
|
||||
|
||||
def assert_video_timestamps_within_bounds(aggr_ds):
|
||||
"""Test that all video timestamps are within valid bounds for their respective video files.
|
||||
|
||||
@@ -260,7 +283,11 @@ def assert_video_timestamps_within_bounds(aggr_ds):
|
||||
|
||||
|
||||
def test_aggregate_datasets(tmp_path, lerobot_dataset_factory):
|
||||
"""Test basic aggregation functionality with standard parameters."""
|
||||
"""Test basic aggregation functionality with standard parameters.
|
||||
|
||||
Source datasets include both RGB and depth video features so the same
|
||||
aggregation flow is exercised on the ``is_depth_map`` branch.
|
||||
"""
|
||||
ds_0_num_frames = 400
|
||||
ds_1_num_frames = 800
|
||||
ds_0_num_episodes = 10
|
||||
@@ -272,14 +299,21 @@ def test_aggregate_datasets(tmp_path, lerobot_dataset_factory):
|
||||
repo_id=f"{DUMMY_REPO_ID}_0",
|
||||
total_episodes=ds_0_num_episodes,
|
||||
total_frames=ds_0_num_frames,
|
||||
camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
|
||||
)
|
||||
ds_1 = lerobot_dataset_factory(
|
||||
root=tmp_path / "test_1",
|
||||
repo_id=f"{DUMMY_REPO_ID}_1",
|
||||
total_episodes=ds_1_num_episodes,
|
||||
total_frames=ds_1_num_frames,
|
||||
camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
|
||||
)
|
||||
|
||||
# Confirm depth was actually wired into the source datasets so the
|
||||
# rest of the assertions exercise the depth aggregation path.
|
||||
assert len(ds_0.meta.depth_keys) > 0, "ds_0 should expose at least one depth key"
|
||||
assert len(ds_1.meta.depth_keys) > 0, "ds_1 should expose at least one depth key"
|
||||
|
||||
aggregate_datasets(
|
||||
repo_ids=[ds_0.repo_id, ds_1.repo_id],
|
||||
roots=[ds_0.root, ds_1.root],
|
||||
@@ -306,6 +340,7 @@ def test_aggregate_datasets(tmp_path, lerobot_dataset_factory):
|
||||
assert_episode_indices_updated_correctly(aggr_ds, ds_0, ds_1)
|
||||
assert_video_frames_integrity(aggr_ds, ds_0, ds_1)
|
||||
assert_video_timestamps_within_bounds(aggr_ds)
|
||||
assert_depth_keys_preserved(aggr_ds, ds_0, ds_1)
|
||||
assert_dataset_iteration_works(aggr_ds)
|
||||
|
||||
|
||||
@@ -423,7 +458,11 @@ def test_aggregate_incomplete_video_encoder_info_warns_and_nuls_encoders(
|
||||
|
||||
|
||||
def test_aggregate_with_low_threshold(tmp_path, lerobot_dataset_factory):
|
||||
"""Test aggregation with small file size limits to force file rotation/sharding."""
|
||||
"""Test aggregation with small file size limits to force file rotation/sharding.
|
||||
|
||||
Depth video features are included to verify that file rotation/concat
|
||||
correctly handles depth-marked features alongside regular RGB ones.
|
||||
"""
|
||||
ds_0_num_episodes = ds_1_num_episodes = 10
|
||||
ds_0_num_frames = ds_1_num_frames = 400
|
||||
|
||||
@@ -432,14 +471,19 @@ def test_aggregate_with_low_threshold(tmp_path, lerobot_dataset_factory):
|
||||
repo_id=f"{DUMMY_REPO_ID}_small_0",
|
||||
total_episodes=ds_0_num_episodes,
|
||||
total_frames=ds_0_num_frames,
|
||||
camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
|
||||
)
|
||||
ds_1 = lerobot_dataset_factory(
|
||||
root=tmp_path / "small_1",
|
||||
repo_id=f"{DUMMY_REPO_ID}_small_1",
|
||||
total_episodes=ds_1_num_episodes,
|
||||
total_frames=ds_1_num_frames,
|
||||
camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
|
||||
)
|
||||
|
||||
assert len(ds_0.meta.depth_keys) > 0, "ds_0 should expose at least one depth key"
|
||||
assert len(ds_1.meta.depth_keys) > 0, "ds_1 should expose at least one depth key"
|
||||
|
||||
# Use the new configurable parameters to force file rotation
|
||||
aggregate_datasets(
|
||||
repo_ids=[ds_0.repo_id, ds_1.repo_id],
|
||||
@@ -470,6 +514,7 @@ def test_aggregate_with_low_threshold(tmp_path, lerobot_dataset_factory):
|
||||
assert_episode_indices_updated_correctly(aggr_ds, ds_0, ds_1)
|
||||
assert_video_frames_integrity(aggr_ds, ds_0, ds_1)
|
||||
assert_video_timestamps_within_bounds(aggr_ds)
|
||||
assert_depth_keys_preserved(aggr_ds, ds_0, ds_1)
|
||||
assert_dataset_iteration_works(aggr_ds)
|
||||
|
||||
# Check that multiple files were actually created due to small size limits
|
||||
@@ -489,7 +534,8 @@ def test_video_timestamps_regression(tmp_path, lerobot_dataset_factory):
|
||||
"""Regression test for video timestamp bug when merging datasets.
|
||||
|
||||
This test specifically checks that video timestamps are correctly calculated
|
||||
and accumulated when merging multiple datasets.
|
||||
and accumulated when merging multiple datasets. Depth video features are
|
||||
included so depth timestamps are also covered by the regression.
|
||||
"""
|
||||
datasets = []
|
||||
for i in range(3):
|
||||
@@ -498,9 +544,13 @@ def test_video_timestamps_regression(tmp_path, lerobot_dataset_factory):
|
||||
repo_id=f"{DUMMY_REPO_ID}_regression_{i}",
|
||||
total_episodes=2,
|
||||
total_frames=100,
|
||||
camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
|
||||
)
|
||||
datasets.append(ds)
|
||||
|
||||
for i, ds in enumerate(datasets):
|
||||
assert len(ds.meta.depth_keys) > 0, f"Dataset {i} should expose at least one depth key"
|
||||
|
||||
aggregate_datasets(
|
||||
repo_ids=[ds.repo_id for ds in datasets],
|
||||
roots=[ds.root for ds in datasets],
|
||||
@@ -517,12 +567,21 @@ def test_video_timestamps_regression(tmp_path, lerobot_dataset_factory):
|
||||
aggr_ds = LeRobotDataset(f"{DUMMY_REPO_ID}_regression_aggr", root=tmp_path / "regression_aggr")
|
||||
|
||||
assert_video_timestamps_within_bounds(aggr_ds)
|
||||
# Depth keys must survive the merge for the regression to cover the
|
||||
# ``is_depth_map`` decoding branch.
|
||||
assert set(aggr_ds.meta.depth_keys) == set(datasets[0].meta.depth_keys)
|
||||
|
||||
depth_keys = set(aggr_ds.meta.depth_keys)
|
||||
for i in range(len(aggr_ds)):
|
||||
item = aggr_ds[i]
|
||||
for key in aggr_ds.meta.video_keys:
|
||||
assert key in item, f"Video key {key} missing from item {i}"
|
||||
assert item[key].shape[0] == 3, f"Expected 3 channels for video key {key}"
|
||||
# Depth frames are single-channel (1, H, W) after dequantization;
|
||||
# standard RGB frames keep the 3-channel layout.
|
||||
expected_channels = 1 if key in depth_keys else 3
|
||||
assert item[key].shape[0] == expected_channels, (
|
||||
f"Expected {expected_channels} channels for video key {key}, got {item[key].shape}"
|
||||
)
|
||||
|
||||
|
||||
def assert_image_schema_preserved(aggr_ds):
|
||||
@@ -639,25 +698,31 @@ def test_aggregate_image_datasets(tmp_path, lerobot_dataset_factory):
|
||||
ds_0_num_episodes = 2
|
||||
ds_1_num_episodes = 3
|
||||
|
||||
# Create two image-based datasets (use_videos=False)
|
||||
# Create two image-based datasets (use_videos=False) with a mix of RGB
|
||||
# and depth-marked cameras so the depth path is exercised in image mode.
|
||||
ds_0 = lerobot_dataset_factory(
|
||||
root=tmp_path / "image_0",
|
||||
repo_id=f"{DUMMY_REPO_ID}_image_0",
|
||||
total_episodes=ds_0_num_episodes,
|
||||
total_frames=ds_0_num_frames,
|
||||
use_videos=False, # Image-based dataset
|
||||
use_videos=False,
|
||||
camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
|
||||
)
|
||||
ds_1 = lerobot_dataset_factory(
|
||||
root=tmp_path / "image_1",
|
||||
repo_id=f"{DUMMY_REPO_ID}_image_1",
|
||||
total_episodes=ds_1_num_episodes,
|
||||
total_frames=ds_1_num_frames,
|
||||
use_videos=False, # Image-based dataset
|
||||
use_videos=False,
|
||||
camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
|
||||
)
|
||||
|
||||
# Verify source datasets have image keys
|
||||
assert len(ds_0.meta.image_keys) > 0, "ds_0 should have image keys"
|
||||
assert len(ds_1.meta.image_keys) > 0, "ds_1 should have image keys"
|
||||
# And that the depth marker actually made it onto an image feature.
|
||||
assert len(ds_0.meta.depth_keys) > 0, "ds_0 should expose at least one depth key"
|
||||
assert len(ds_1.meta.depth_keys) > 0, "ds_1 should expose at least one depth key"
|
||||
|
||||
# Aggregate the datasets
|
||||
aggregate_datasets(
|
||||
@@ -692,6 +757,7 @@ def test_aggregate_image_datasets(tmp_path, lerobot_dataset_factory):
|
||||
# Image-specific assertions
|
||||
assert_image_schema_preserved(aggr_ds)
|
||||
assert_image_frames_integrity(aggr_ds, ds_0, ds_1)
|
||||
assert_depth_keys_preserved(aggr_ds, ds_0, ds_1)
|
||||
|
||||
# Verify images can be accessed and have correct shape
|
||||
sample_item = aggr_ds[0]
|
||||
|
||||
@@ -35,7 +35,11 @@ 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)
|
||||
is_depth = "depth" in str(path)
|
||||
channels = 1 if is_depth else 3
|
||||
out_dtype = np.uint16 if is_depth else dtype
|
||||
arr = np.arange(channels * 32 * 32, dtype=out_dtype).reshape(channels, 32, 32)
|
||||
return arr if channel_first else arr.transpose(1, 2, 0)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -168,22 +172,33 @@ def test_get_feature_stats_single_value():
|
||||
|
||||
|
||||
def test_compute_episode_stats():
|
||||
depth_key = "observation.images.depth"
|
||||
episode_data = {
|
||||
OBS_IMAGE: [f"image_{i}.jpg" for i in range(100)],
|
||||
depth_key: [f"depth_{i}.tiff" for i in range(100)],
|
||||
OBS_STATE: np.random.rand(100, 10),
|
||||
}
|
||||
features = {
|
||||
OBS_IMAGE: {"dtype": "image"},
|
||||
depth_key: {"dtype": "image", "info": {"is_depth_map": True}},
|
||||
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 OBS_IMAGE in stats and depth_key in stats and OBS_STATE in stats
|
||||
assert stats[OBS_IMAGE]["count"].item() == 100
|
||||
assert stats[depth_key]["count"].item() == 100
|
||||
assert stats[OBS_STATE]["count"].item() == 100
|
||||
assert stats[OBS_IMAGE]["mean"].shape == (3, 1, 1)
|
||||
assert stats[depth_key]["mean"].shape == (1, 1, 1)
|
||||
# Depth keeps raw values: max far exceeds 255, proving no /255 and no uint8 downcast.
|
||||
assert stats[depth_key]["min"].item() == 0.0
|
||||
assert stats[depth_key]["max"].item() == 1023.0
|
||||
# RGB is normalized to [0, 1].
|
||||
np.testing.assert_allclose(stats[OBS_IMAGE]["min"], 0.0)
|
||||
np.testing.assert_allclose(stats[OBS_IMAGE]["max"], 1.0)
|
||||
|
||||
|
||||
def test_assert_type_and_shape_valid():
|
||||
@@ -618,25 +633,31 @@ def test_compute_episode_stats_with_custom_quantiles():
|
||||
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)]
|
||||
depth_paths = [f"depth_{i}.tiff" for i in range(50)]
|
||||
episode_data = {
|
||||
"observation.image": image_paths,
|
||||
"observation.images.depth": depth_paths,
|
||||
"action": np.random.normal(0, 1, (50, 5)),
|
||||
}
|
||||
features = {
|
||||
"observation.image": {"dtype": "image"},
|
||||
"observation.images.depth": {"dtype": "image", "info": {"is_depth_map": True}},
|
||||
"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)
|
||||
# RGB image quantiles should be normalized and per-channel.
|
||||
for q in ("q01", "q50", "q99"):
|
||||
assert stats["observation.image"][q].shape == (3, 1, 1)
|
||||
|
||||
# Depth quantiles are single-channel and kept in raw (un-normalized) units.
|
||||
for q in ("q01", "q50", "q99"):
|
||||
assert stats["observation.images.depth"][q].shape == (1, 1, 1)
|
||||
# Depth max stays in raw units (not /255, not uint8-capped); RGB is normalized.
|
||||
assert stats["observation.images.depth"]["max"].item() == 1023.0
|
||||
np.testing.assert_allclose(stats["observation.image"]["max"], 1.0)
|
||||
|
||||
# Action quantiles should have correct shape
|
||||
assert stats["action"]["q01"].shape == (5,)
|
||||
|
||||
@@ -59,11 +59,13 @@ def _make_dummy_stats(features: dict) -> dict:
|
||||
stats = {}
|
||||
for key, ft in features.items():
|
||||
if ft["dtype"] in ("image", "video"):
|
||||
channels = ft["shape"][-1]
|
||||
stat_shape = (channels, 1, 1)
|
||||
stats[key] = {
|
||||
"max": np.ones((3, 1, 1), dtype=np.float32),
|
||||
"mean": np.full((3, 1, 1), 0.5, dtype=np.float32),
|
||||
"min": np.zeros((3, 1, 1), dtype=np.float32),
|
||||
"std": np.full((3, 1, 1), 0.25, dtype=np.float32),
|
||||
"max": np.ones(stat_shape, dtype=np.float32),
|
||||
"mean": np.full(stat_shape, 0.5, dtype=np.float32),
|
||||
"min": np.zeros(stat_shape, dtype=np.float32),
|
||||
"std": np.full(stat_shape, 0.25, dtype=np.float32),
|
||||
"count": np.array([5]),
|
||||
}
|
||||
elif ft["dtype"] in ("float32", "float64", "int64"):
|
||||
@@ -142,6 +144,45 @@ def test_create_without_videos_has_no_video_path(tmp_path):
|
||||
assert meta.video_keys == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("marker_field", "marker_key"),
|
||||
[
|
||||
("info", "is_depth_map"),
|
||||
("info", "video.is_depth_map"),
|
||||
("video_info", "video.is_depth_map"),
|
||||
],
|
||||
ids=["info.is_depth_map", "info.video.is_depth_map_legacy", "video_info.video.is_depth_map_legacy"],
|
||||
)
|
||||
def test_depth_keys_property_filters_by_marker(tmp_path, marker_field, marker_key):
|
||||
"""``depth_keys`` recognises the canonical and the two legacy marker variants."""
|
||||
depth_feature = {
|
||||
"dtype": "video",
|
||||
"shape": (64, 96, 1),
|
||||
"names": ["height", "width", "channels"],
|
||||
marker_field: {marker_key: True},
|
||||
}
|
||||
features = {
|
||||
**VIDEO_FEATURES,
|
||||
"observation.images.laptop_depth": depth_feature,
|
||||
}
|
||||
meta = LeRobotDatasetMetadata.create(
|
||||
repo_id="test/depth_keys",
|
||||
fps=DEFAULT_FPS,
|
||||
features=features,
|
||||
root=tmp_path / f"depth_keys_{marker_field}_{marker_key.replace('.', '_')}",
|
||||
)
|
||||
|
||||
assert set(meta.video_keys) == {"observation.images.laptop", "observation.images.laptop_depth"}
|
||||
assert meta.depth_keys == ["observation.images.laptop_depth"]
|
||||
|
||||
|
||||
def test_depth_keys_empty_when_no_marker(tmp_path):
|
||||
meta = LeRobotDatasetMetadata.create(
|
||||
repo_id="test/no_depth", fps=DEFAULT_FPS, features=VIDEO_FEATURES, root=tmp_path / "no_depth"
|
||||
)
|
||||
assert meta.depth_keys == []
|
||||
|
||||
|
||||
def test_create_raises_on_existing_directory(tmp_path):
|
||||
"""create() raises if root directory already exists."""
|
||||
root = tmp_path / "existing"
|
||||
|
||||
@@ -24,7 +24,7 @@ import torch
|
||||
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
||||
|
||||
|
||||
from lerobot.configs import VideoEncoderConfig
|
||||
from lerobot.configs import DepthEncoderConfig, RGBEncoderConfig
|
||||
from lerobot.datasets.dataset_tools import (
|
||||
add_features,
|
||||
convert_image_to_video_dataset,
|
||||
@@ -37,7 +37,9 @@ from lerobot.datasets.dataset_tools import (
|
||||
split_dataset,
|
||||
)
|
||||
from lerobot.datasets.io_utils import load_info
|
||||
from tests.datasets.test_video_encoding import _add_frames, require_h264, require_libsvtav1
|
||||
from tests.datasets.test_video_encoding import require_h264, require_hevc, require_libsvtav1
|
||||
from tests.fixtures.constants import DUMMY_DEPTH_FEATURES, DUMMY_DEPTH_KEY
|
||||
from tests.fixtures.dataset_factories import add_frames
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -1251,7 +1253,7 @@ def test_convert_image_to_video_dataset(tmp_path):
|
||||
dataset=source_dataset,
|
||||
output_dir=output_dir,
|
||||
repo_id="lerobot/pusht_video",
|
||||
camera_encoder=VideoEncoderConfig(
|
||||
rgb_encoder=RGBEncoderConfig(
|
||||
vcodec="libsvtav1",
|
||||
pix_fmt="yuv420p",
|
||||
g=2,
|
||||
@@ -1332,9 +1334,131 @@ def test_convert_image_to_video_dataset_subset_episodes(tmp_path):
|
||||
shutil.rmtree(output_dir)
|
||||
|
||||
|
||||
@require_libsvtav1
|
||||
@require_hevc
|
||||
def test_convert_image_to_video_dataset_depth(tmp_path, empty_lerobot_dataset_factory):
|
||||
"""Depth image features convert to depth videos using the depth encoder.
|
||||
|
||||
Mirrors :func:`test_convert_image_to_video_dataset` but with a small local
|
||||
image dataset that mixes an RGB camera with a depth camera, so the
|
||||
``depth_keys`` → ``depth_encoder`` routing and ``is_depth_map`` preservation
|
||||
are exercised end-to-end.
|
||||
"""
|
||||
features = {
|
||||
"action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]},
|
||||
"observation.images.cam": {
|
||||
"dtype": "image",
|
||||
"shape": (64, 96, 3),
|
||||
"names": ["height", "width", "channels"],
|
||||
},
|
||||
"observation.images.depth": {
|
||||
"dtype": "image",
|
||||
"shape": (64, 96, 1),
|
||||
"names": ["height", "width", "channels"],
|
||||
"info": {"is_depth_map": True},
|
||||
},
|
||||
}
|
||||
source_dataset = empty_lerobot_dataset_factory(
|
||||
root=tmp_path / "img_ds",
|
||||
features=features,
|
||||
use_videos=False,
|
||||
)
|
||||
|
||||
add_frames(source_dataset, num_frames=4)
|
||||
source_dataset.save_episode()
|
||||
source_dataset.finalize()
|
||||
|
||||
# Source is an image dataset with the depth marker on the depth camera.
|
||||
assert len(source_dataset.meta.video_keys) == 0
|
||||
assert "observation.images.depth" in source_dataset.meta.depth_keys
|
||||
|
||||
output_dir = tmp_path / "video_ds"
|
||||
with (
|
||||
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
|
||||
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
|
||||
):
|
||||
mock_get_safe_version.return_value = "v3.0"
|
||||
mock_snapshot_download.return_value = str(output_dir)
|
||||
|
||||
# Use non-default quantization params so the persisted metadata must
|
||||
# come from the depth encoder (not RGB encoder defaults).
|
||||
depth_encoder = DepthEncoderConfig(
|
||||
vcodec="hevc",
|
||||
pix_fmt="gray12le",
|
||||
g=2,
|
||||
crf=30,
|
||||
depth_min=0.05,
|
||||
depth_max=8.0,
|
||||
shift=2.0,
|
||||
use_log=False,
|
||||
)
|
||||
video_dataset = convert_image_to_video_dataset(
|
||||
dataset=source_dataset,
|
||||
output_dir=output_dir,
|
||||
repo_id="dummy/depth_video",
|
||||
rgb_encoder=RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
|
||||
depth_encoder=depth_encoder,
|
||||
num_workers=1,
|
||||
)
|
||||
|
||||
# Both cameras are now videos, and the depth marker survived the conversion.
|
||||
assert "observation.images.cam" in video_dataset.meta.video_keys
|
||||
assert "observation.images.depth" in video_dataset.meta.video_keys
|
||||
assert "observation.images.depth" in video_dataset.meta.depth_keys
|
||||
assert "observation.images.cam" not in video_dataset.meta.depth_keys
|
||||
|
||||
depth_path = video_dataset.root / video_dataset.meta.get_video_file_path(0, "observation.images.depth")
|
||||
assert depth_path.exists(), f"Depth video file should exist: {depth_path}"
|
||||
|
||||
# The persisted depth-video metadata must carry the depth quantization params
|
||||
# from the depth encoder (so frames dequantize correctly on read), and the RGB
|
||||
# camera must not be marked as a depth map.
|
||||
persisted_info = load_info(video_dataset.root)
|
||||
depth_info = persisted_info.features["observation.images.depth"]["info"]
|
||||
assert depth_info["is_depth_map"] is True
|
||||
assert DepthEncoderConfig.from_video_info(depth_info) == depth_encoder
|
||||
|
||||
cam_info = persisted_info.features["observation.images.cam"]["info"]
|
||||
assert cam_info.get("is_depth_map") is False
|
||||
assert "video.codec" in cam_info
|
||||
|
||||
|
||||
# ─── reencode_dataset ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
@require_hevc
|
||||
def test_reencode_dataset_depth_uses_depth_encoder(tmp_path, empty_lerobot_dataset_factory):
|
||||
"""Depth videos are re-encoded with the depth encoder and keep their depth metadata.
|
||||
|
||||
Depth-focused companion to :func:`test_reencode_dataset_multi_key_multiprocessing`.
|
||||
"""
|
||||
initial_cfg = DepthEncoderConfig(vcodec="hevc", pix_fmt="gray12le", g=2, crf=30)
|
||||
dataset = empty_lerobot_dataset_factory(
|
||||
root=tmp_path / "ds",
|
||||
features=DUMMY_DEPTH_FEATURES,
|
||||
use_videos=True,
|
||||
depth_encoder=initial_cfg,
|
||||
)
|
||||
|
||||
add_frames(dataset, num_frames=4)
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
|
||||
assert DUMMY_DEPTH_KEY in dataset.meta.depth_keys
|
||||
|
||||
target_cfg = DepthEncoderConfig(vcodec="hevc", pix_fmt="gray12le", g=6, crf=23)
|
||||
result = reencode_dataset(dataset, depth_encoder=target_cfg, num_workers=0)
|
||||
|
||||
assert result is dataset
|
||||
|
||||
persisted_info = load_info(dataset.root)
|
||||
depth_info = persisted_info.features[DUMMY_DEPTH_KEY].get("info", {})
|
||||
# Re-encode applied the new codec parameters to the depth video ...
|
||||
assert DepthEncoderConfig.from_video_info(depth_info) == target_cfg
|
||||
# ... while preserving the depth marker.
|
||||
assert depth_info["is_depth_map"] is True
|
||||
|
||||
|
||||
@require_libsvtav1
|
||||
@require_h264
|
||||
def test_reencode_dataset_multi_key_multiprocessing(
|
||||
@@ -1342,29 +1466,29 @@ def test_reencode_dataset_multi_key_multiprocessing(
|
||||
):
|
||||
"""Re-encode a two-camera dataset with num_workers=2 and verify metadata refresh."""
|
||||
features = features_factory(use_videos=True)
|
||||
initial_cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
|
||||
initial_cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
|
||||
dataset = empty_lerobot_dataset_factory(
|
||||
root=tmp_path / "ds",
|
||||
features=features,
|
||||
use_videos=True,
|
||||
camera_encoder=initial_cfg,
|
||||
rgb_encoder=initial_cfg,
|
||||
)
|
||||
|
||||
_add_frames(dataset, num_frames=4)
|
||||
add_frames(dataset, num_frames=4)
|
||||
dataset.save_episode()
|
||||
_add_frames(dataset, num_frames=4)
|
||||
add_frames(dataset, num_frames=4)
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
|
||||
assert len(dataset.meta.video_keys) == 2
|
||||
|
||||
target_cfg = VideoEncoderConfig(vcodec="h264", g=6, crf=23, pix_fmt="yuv420p")
|
||||
target_cfg = RGBEncoderConfig(vcodec="h264", g=6, crf=23, pix_fmt="yuv420p")
|
||||
|
||||
result = reencode_dataset(dataset, camera_encoder=target_cfg, num_workers=2)
|
||||
result = reencode_dataset(dataset, rgb_encoder=target_cfg, num_workers=2)
|
||||
|
||||
assert result is dataset
|
||||
|
||||
persisted_info = load_info(dataset.root)
|
||||
for vk in dataset.meta.video_keys:
|
||||
persisted_encoder = VideoEncoderConfig.from_video_info(persisted_info.features[vk].get("info", {}))
|
||||
persisted_encoder = RGBEncoderConfig.from_video_info(persisted_info.features[vk].get("info", {}))
|
||||
assert persisted_encoder == target_cfg
|
||||
|
||||
@@ -53,8 +53,8 @@ def _make_frame(features: dict, task: str = "Dummy task") -> dict:
|
||||
# ── Existing encode_video_worker tests ───────────────────────────────
|
||||
|
||||
|
||||
def test_encode_video_worker_forwards_camera_encoder(tmp_path):
|
||||
"""_encode_video_worker forwards camera_encoder to encode_video_frames."""
|
||||
def test_encode_video_worker_forwards_video_encoder(tmp_path):
|
||||
"""_encode_video_worker forwards video_encoder to encode_video_frames."""
|
||||
video_key = "observation.images.laptop"
|
||||
fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0)
|
||||
img_dir = tmp_path / Path(fpath).parent
|
||||
@@ -74,16 +74,16 @@ def test_encode_video_worker_forwards_camera_encoder(tmp_path):
|
||||
0,
|
||||
tmp_path,
|
||||
fps=30,
|
||||
camera_encoder=VideoEncoderConfig(vcodec="h264", preset=None),
|
||||
video_encoder=VideoEncoderConfig(vcodec="h264", preset=None),
|
||||
encoder_threads=4,
|
||||
)
|
||||
|
||||
assert captured_kwargs["camera_encoder"].vcodec == "h264"
|
||||
assert captured_kwargs["video_encoder"].vcodec == "h264"
|
||||
assert captured_kwargs["encoder_threads"] == 4
|
||||
|
||||
|
||||
def test_encode_video_worker_default_camera_encoder(tmp_path):
|
||||
"""_encode_video_worker passes None camera_encoder which encode_video_frames defaults."""
|
||||
def test_encode_video_worker_default_video_encoder(tmp_path):
|
||||
"""_encode_video_worker passes None video_encoder which encode_video_frames defaults."""
|
||||
video_key = "observation.images.laptop"
|
||||
fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0)
|
||||
img_dir = tmp_path / Path(fpath).parent
|
||||
@@ -100,7 +100,7 @@ def test_encode_video_worker_default_camera_encoder(tmp_path):
|
||||
with patch("lerobot.datasets.dataset_writer.encode_video_frames", side_effect=mock_encode):
|
||||
_encode_video_worker(video_key, 0, tmp_path, fps=30)
|
||||
|
||||
assert captured_kwargs["camera_encoder"] is None
|
||||
assert captured_kwargs["video_encoder"] is None
|
||||
assert captured_kwargs["encoder_threads"] is None
|
||||
|
||||
|
||||
|
||||
@@ -1534,6 +1534,10 @@ def test_valid_video_codecs_constant():
|
||||
assert "auto" in VALID_VIDEO_CODECS
|
||||
assert "h264_videotoolbox" in VALID_VIDEO_CODECS
|
||||
assert "h264_nvenc" in VALID_VIDEO_CODECS
|
||||
assert "h264_vaapi" in VALID_VIDEO_CODECS
|
||||
assert "h264_qsv" in VALID_VIDEO_CODECS
|
||||
assert "hevc_videotoolbox" in VALID_VIDEO_CODECS
|
||||
assert "hevc_nvenc" in VALID_VIDEO_CODECS
|
||||
assert len(VALID_VIDEO_CODECS) == 10
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,247 @@
|
||||
"""Tests for the depth-integration feature.
|
||||
|
||||
Covers:
|
||||
- ``depth_utils`` quantize/dequantize round-trips and backend agreement.
|
||||
- Image-writer support for single-channel depth.
|
||||
- Hardware-feature → depth flag routing.
|
||||
- Feature-to-file-format routing through the dataset writer.
|
||||
|
||||
Depth metadata detection on ``LeRobotDatasetMetadata.depth_keys`` lives in
|
||||
``test_dataset_metadata.py``. Depth video encoding/decoding lives in
|
||||
``test_video_encoding.py``.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("av", reason="av is required (install lerobot[dataset])")
|
||||
|
||||
import av
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import torch
|
||||
|
||||
from lerobot.configs import DepthEncoderConfig
|
||||
from lerobot.configs.video import (
|
||||
DEFAULT_DEPTH_MAX,
|
||||
DEFAULT_DEPTH_MIN,
|
||||
DEPTH_METER_UNIT,
|
||||
DEPTH_MILLIMETER_UNIT,
|
||||
DEPTH_QMAX,
|
||||
)
|
||||
from lerobot.datasets.depth_utils import dequantize_depth, quantize_depth
|
||||
from lerobot.datasets.image_writer import image_array_to_pil_image, write_image
|
||||
from tests.fixtures.constants import (
|
||||
DEFAULT_FPS,
|
||||
DUMMY_CAMERA_FEATURES,
|
||||
DUMMY_CAMERA_FEATURES_WITH_DEPTH,
|
||||
DUMMY_CHW,
|
||||
DUMMY_DEPTH_CAMERA_FEATURES,
|
||||
DUMMY_REPO_ID,
|
||||
)
|
||||
from tests.fixtures.dataset_factories import add_frames
|
||||
|
||||
_, H, W = DUMMY_CHW
|
||||
|
||||
|
||||
def _depth_metres_ramp() -> np.ndarray:
|
||||
"""Linearly-spaced float32 depth in metres covering the default range."""
|
||||
return np.linspace(DEFAULT_DEPTH_MIN, DEFAULT_DEPTH_MAX, H * W, dtype=np.float32).reshape(H, W)
|
||||
|
||||
|
||||
# ── 1. Quantize / dequantize round-trips ──────────────────────────────
|
||||
|
||||
|
||||
class TestQuantizeDequantize:
|
||||
"""Numerical contract of ``quantize_depth`` / ``dequantize_depth``."""
|
||||
|
||||
@pytest.mark.parametrize("use_log", [False, True])
|
||||
@pytest.mark.parametrize("output_unit", [DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT])
|
||||
@pytest.mark.parametrize("output_channel_last", [False, True])
|
||||
def test_roundtrip(self, use_log, output_unit, output_channel_last):
|
||||
"""quantize → dequantize recovers depth; layout and unit are honored."""
|
||||
depth = _depth_metres_ramp()
|
||||
quantized = quantize_depth(depth, use_log=use_log, video_backend=None)
|
||||
recovered = dequantize_depth(
|
||||
quantized,
|
||||
use_log=use_log,
|
||||
output_unit=output_unit,
|
||||
output_tensor=False,
|
||||
output_channel_last=output_channel_last,
|
||||
)
|
||||
|
||||
expected_shape = (H, W, 1) if output_channel_last else (1, H, W)
|
||||
assert recovered.shape == expected_shape
|
||||
|
||||
recovered_m = recovered.astype(np.float32)
|
||||
if output_unit == DEPTH_MILLIMETER_UNIT:
|
||||
recovered_m = recovered_m / 1000.0
|
||||
recovered_2d = recovered_m[..., 0] if output_channel_last else recovered_m[0]
|
||||
|
||||
if use_log:
|
||||
# Log mode: tighter near-range error than far-range (the whole point).
|
||||
near = depth < 1.0
|
||||
far = depth > 8.0
|
||||
err_near = np.abs(recovered_2d[near] - depth[near])
|
||||
err_far = np.abs(recovered_2d[far] - depth[far])
|
||||
assert err_near.mean() < err_far.mean()
|
||||
else:
|
||||
# Linear mode: bounded by quant step + 1 mm of unit-conversion rounding.
|
||||
tol = (DEFAULT_DEPTH_MAX - DEFAULT_DEPTH_MIN) / DEPTH_QMAX + 1e-3
|
||||
np.testing.assert_allclose(recovered_2d, depth, atol=tol)
|
||||
|
||||
@pytest.mark.parametrize("use_log", [False, True])
|
||||
@pytest.mark.parametrize("output_unit", [DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT])
|
||||
def test_numpy_torch_agree(self, use_log, output_unit):
|
||||
"""Batched torch path produces the same values as the numpy path."""
|
||||
batch_size = 3
|
||||
per_frame = np.linspace(0, DEPTH_QMAX, H * W, dtype=np.uint16).reshape(H, W)
|
||||
batch_np = np.broadcast_to(per_frame[None, None, ...], (batch_size, 1, H, W)).copy()
|
||||
batch_t = torch.from_numpy(batch_np.astype(np.int32)) # torch.uint16 support is patchy.
|
||||
|
||||
ref = dequantize_depth(batch_np, use_log=use_log, output_unit=output_unit, output_tensor=False)
|
||||
out = dequantize_depth(batch_t, use_log=use_log, output_unit=output_unit, output_tensor=True)
|
||||
|
||||
assert isinstance(out, torch.Tensor)
|
||||
assert out.shape == (batch_size, 1, H, W)
|
||||
# ``m``: float32 noise (~10 µm in log mode, after ``exp``) — still 200× below the ~2 mm quant step.
|
||||
# ``mm`` + tensor stays in float32 (no uint16 round-trip), so allow 1 mm slop.
|
||||
atol = 1e-5 if output_unit == DEPTH_METER_UNIT else 1.0
|
||||
np.testing.assert_allclose(out.cpu().numpy().astype(np.float64), ref.astype(np.float64), atol=atol)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"input_shape,output_shape",
|
||||
[
|
||||
((H, W), (1, H, W)),
|
||||
((1, H, W), (1, H, W)),
|
||||
((H, W, 1), (1, H, W)),
|
||||
((3, 1, H, W), (3, 1, H, W)),
|
||||
((3, H, W, 1), (3, 1, H, W)),
|
||||
],
|
||||
)
|
||||
def test_input_layouts_accepted(self, input_shape, output_shape):
|
||||
"""All documented input layouts decode to the channel-first default."""
|
||||
quantized = np.full(input_shape, DEPTH_QMAX // 2, dtype=np.uint16)
|
||||
out = dequantize_depth(quantized, output_unit=DEPTH_METER_UNIT, output_tensor=False)
|
||||
assert out.shape == output_shape
|
||||
|
||||
def test_pyav_frame_roundtrip(self):
|
||||
"""quantize → av.VideoFrame → dequantize works."""
|
||||
depth = _depth_metres_ramp()
|
||||
frame = quantize_depth(depth, use_log=False, video_backend="pyav")
|
||||
assert isinstance(frame, av.VideoFrame)
|
||||
|
||||
recovered = dequantize_depth(frame, use_log=False, output_unit=DEPTH_METER_UNIT, output_tensor=False)
|
||||
assert recovered.shape == (1, H, W)
|
||||
tol = (DEFAULT_DEPTH_MAX - DEFAULT_DEPTH_MIN) / DEPTH_QMAX + 1e-3
|
||||
np.testing.assert_allclose(recovered[0], depth, atol=tol)
|
||||
|
||||
def test_invalid_log_params_raises(self):
|
||||
with pytest.raises(ValueError, match=r"depth_min \+ shift must be positive"):
|
||||
quantize_depth(_depth_metres_ramp(), depth_min=1.0, shift=-2.0, use_log=True, video_backend=None)
|
||||
|
||||
|
||||
# ── 2. Image writer depth support ─────────────────────────────────────
|
||||
|
||||
|
||||
class TestImageWriterDepth:
|
||||
"""``image_array_to_pil_image`` and ``write_image`` for depth maps."""
|
||||
|
||||
@pytest.mark.parametrize("dtype,expected_mode", [(np.uint16, "I;16"), (np.float32, "F")])
|
||||
@pytest.mark.parametrize("shape", [(H, W), (H, W, 1), (1, H, W)])
|
||||
def test_pil_depth_modes_and_squeeze(self, dtype, expected_mode, shape):
|
||||
"""Single-channel depth converts to PIL with the right mode and (W, H) size."""
|
||||
arr = np.zeros(shape, dtype=dtype)
|
||||
img = image_array_to_pil_image(arr)
|
||||
assert img.mode == expected_mode
|
||||
assert img.size == (W, H)
|
||||
|
||||
def test_write_image_tiff_roundtrip(self, tmp_path):
|
||||
"""uint16 depth round-trips through .tiff."""
|
||||
arr = np.arange(H * W, dtype=np.uint16).reshape(H, W)
|
||||
fpath = tmp_path / "depth.tiff"
|
||||
write_image(arr, fpath)
|
||||
with PIL.Image.open(fpath) as loaded:
|
||||
recovered = np.array(loaded)
|
||||
np.testing.assert_array_equal(recovered, arr)
|
||||
|
||||
|
||||
# ── 3. Hardware-feature → depth flag ──────────────────────────────────
|
||||
|
||||
|
||||
class TestHwToDatasetFeaturesDepth:
|
||||
"""``hw_to_dataset_features`` flags single-channel cameras as depth."""
|
||||
|
||||
@pytest.mark.parametrize("channels,is_depth", [(1, True), (3, False)])
|
||||
def test_depth_marker_by_channels(self, channels, is_depth):
|
||||
from lerobot.utils.feature_utils import hw_to_dataset_features
|
||||
|
||||
features = hw_to_dataset_features({"cam": (480, 640, channels)}, prefix="observation")
|
||||
assert features["observation.images.cam"]["info"]["is_depth_map"] is is_depth
|
||||
|
||||
def test_invalid_channel_count_raises(self):
|
||||
from lerobot.utils.feature_utils import hw_to_dataset_features
|
||||
|
||||
with pytest.raises(ValueError, match="Expected a 3-tuple"):
|
||||
hw_to_dataset_features({"cam": (480, 640, 2)}, prefix="observation")
|
||||
|
||||
|
||||
# ── 4. Feature-to-file-format routing ────────────────────────────────
|
||||
|
||||
|
||||
# Keys derived from DUMMY_CAMERA_FEATURES_WITH_DEPTH; pick one RGB and the depth camera.
|
||||
RGB_KEY = next(iter(DUMMY_CAMERA_FEATURES))
|
||||
DEPTH_KEY = next(iter(DUMMY_DEPTH_CAMERA_FEATURES))
|
||||
|
||||
|
||||
class TestFeatureFileRouting:
|
||||
"""Depth vs RGB features route to the correct file format."""
|
||||
|
||||
NUM_FRAMES = 5
|
||||
|
||||
def test_image_mode_depth_tiff_rgb_png(self, tmp_path, features_factory):
|
||||
"""Without video encoding: depth → .tiff, RGB → .png."""
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
|
||||
features = features_factory(camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, use_videos=False)
|
||||
dataset = LeRobotDataset.create(
|
||||
repo_id=DUMMY_REPO_ID,
|
||||
fps=DEFAULT_FPS,
|
||||
features=features,
|
||||
root=tmp_path / "ds",
|
||||
use_videos=False,
|
||||
)
|
||||
|
||||
add_frames(dataset, num_frames=self.NUM_FRAMES)
|
||||
|
||||
buf = dataset.writer.episode_buffer
|
||||
assert all(Path(p).suffix == ".tiff" for p in buf[DEPTH_KEY])
|
||||
assert all(Path(p).suffix == ".png" for p in buf[RGB_KEY])
|
||||
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
|
||||
def test_video_mode_depth_uses_depth_encoder(self, tmp_path, features_factory):
|
||||
"""With streaming video encoding: depth → DepthEncoderConfig, RGB does not."""
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
|
||||
features = features_factory(camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, use_videos=True)
|
||||
dataset = LeRobotDataset.create(
|
||||
repo_id=DUMMY_REPO_ID,
|
||||
fps=DEFAULT_FPS,
|
||||
features=features,
|
||||
root=tmp_path / "ds",
|
||||
use_videos=True,
|
||||
streaming_encoding=True,
|
||||
)
|
||||
|
||||
add_frames(dataset, num_frames=self.NUM_FRAMES)
|
||||
|
||||
encoder = dataset.writer._streaming_encoder
|
||||
assert encoder is not None
|
||||
assert isinstance(encoder._threads[DEPTH_KEY].video_encoder, DepthEncoderConfig)
|
||||
assert not isinstance(encoder._threads[RGB_KEY].video_encoder, DepthEncoderConfig)
|
||||
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
@@ -94,7 +94,7 @@ def test_image_array_to_pil_image_pytorch_format(img_array_factory):
|
||||
|
||||
def test_image_array_to_pil_image_single_channel(img_array_factory):
|
||||
img_array = img_array_factory(channels=1)
|
||||
with pytest.raises(NotImplementedError):
|
||||
with pytest.raises(ValueError, match="Unsupported single-channel image dtype"):
|
||||
image_array_to_pil_image(img_array)
|
||||
|
||||
|
||||
@@ -344,7 +344,7 @@ def test_with_different_image_formats(tmp_path, img_array_factory):
|
||||
writer = AsyncImageWriter()
|
||||
try:
|
||||
image_array = img_array_factory()
|
||||
formats = ["png", "jpeg", "bmp"]
|
||||
formats = ["png", "tiff", "tif"]
|
||||
for fmt in formats:
|
||||
fpath = tmp_path / f"test_image.{fmt}"
|
||||
write_image(image_array, fpath)
|
||||
|
||||
@@ -26,7 +26,7 @@ pytest.importorskip("av", reason="av is required (install lerobot[dataset])")
|
||||
|
||||
import av # noqa: E402
|
||||
|
||||
from lerobot.configs import VideoEncoderConfig
|
||||
from lerobot.configs import RGBEncoderConfig
|
||||
from lerobot.datasets.pyav_utils import get_codec
|
||||
from lerobot.datasets.video_utils import (
|
||||
StreamingVideoEncoder,
|
||||
@@ -57,13 +57,11 @@ class TestCameraEncoderThread:
|
||||
result_queue: queue.Queue = queue.Queue(maxsize=1)
|
||||
stop_event = threading.Event()
|
||||
|
||||
enc_cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13)
|
||||
enc_cfg = RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13)
|
||||
encoder_thread = _CameraEncoderThread(
|
||||
video_path=video_path,
|
||||
fps=fps,
|
||||
vcodec=enc_cfg.vcodec,
|
||||
pix_fmt=enc_cfg.pix_fmt,
|
||||
codec_options=enc_cfg.get_codec_options(as_strings=True),
|
||||
video_encoder=enc_cfg,
|
||||
frame_queue=frame_queue,
|
||||
result_queue=result_queue,
|
||||
stop_event=stop_event,
|
||||
@@ -108,13 +106,11 @@ class TestCameraEncoderThread:
|
||||
result_queue: queue.Queue = queue.Queue(maxsize=1)
|
||||
stop_event = threading.Event()
|
||||
|
||||
enc_cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13)
|
||||
enc_cfg = RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13)
|
||||
encoder_thread = _CameraEncoderThread(
|
||||
video_path=video_path,
|
||||
fps=fps,
|
||||
vcodec=enc_cfg.vcodec,
|
||||
pix_fmt=enc_cfg.pix_fmt,
|
||||
codec_options=enc_cfg.get_codec_options(as_strings=True),
|
||||
video_encoder=enc_cfg,
|
||||
frame_queue=frame_queue,
|
||||
result_queue=result_queue,
|
||||
stop_event=stop_event,
|
||||
@@ -142,13 +138,11 @@ class TestCameraEncoderThread:
|
||||
result_queue: queue.Queue = queue.Queue(maxsize=1)
|
||||
stop_event = threading.Event()
|
||||
|
||||
enc_cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13)
|
||||
enc_cfg = RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13)
|
||||
encoder_thread = _CameraEncoderThread(
|
||||
video_path=video_path,
|
||||
fps=fps,
|
||||
vcodec=enc_cfg.vcodec,
|
||||
pix_fmt=enc_cfg.pix_fmt,
|
||||
codec_options=enc_cfg.get_codec_options(as_strings=True),
|
||||
video_encoder=enc_cfg,
|
||||
frame_queue=frame_queue,
|
||||
result_queue=result_queue,
|
||||
stop_event=stop_event,
|
||||
@@ -171,15 +165,15 @@ class TestCameraEncoderThread:
|
||||
|
||||
class TestStreamingVideoEncoder:
|
||||
def _make_encoder_config(self, **kwargs):
|
||||
"""Helper to build a VideoEncoderConfig."""
|
||||
return VideoEncoderConfig(**kwargs)
|
||||
"""Helper to build an RGBEncoderConfig."""
|
||||
return RGBEncoderConfig(**kwargs)
|
||||
|
||||
def test_single_camera_episode(self, tmp_path):
|
||||
"""Test encoding a single camera episode."""
|
||||
video_keys = [f"{OBS_IMAGES}.laptop"]
|
||||
encoder = StreamingVideoEncoder(
|
||||
fps=30,
|
||||
camera_encoder=self._make_encoder_config(
|
||||
rgb_encoder=self._make_encoder_config(
|
||||
vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13
|
||||
),
|
||||
)
|
||||
@@ -211,7 +205,7 @@ class TestStreamingVideoEncoder:
|
||||
video_keys = [f"{OBS_IMAGES}.laptop", f"{OBS_IMAGES}.phone"]
|
||||
encoder = StreamingVideoEncoder(
|
||||
fps=30,
|
||||
camera_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
|
||||
rgb_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
|
||||
)
|
||||
encoder.start_episode(video_keys, tmp_path)
|
||||
|
||||
@@ -237,7 +231,7 @@ class TestStreamingVideoEncoder:
|
||||
video_keys = [f"{OBS_IMAGES}.cam"]
|
||||
encoder = StreamingVideoEncoder(
|
||||
fps=30,
|
||||
camera_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
|
||||
rgb_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
|
||||
)
|
||||
|
||||
for ep in range(3):
|
||||
@@ -263,7 +257,7 @@ class TestStreamingVideoEncoder:
|
||||
video_keys = [f"{OBS_IMAGES}.cam"]
|
||||
encoder = StreamingVideoEncoder(
|
||||
fps=30,
|
||||
camera_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
|
||||
rgb_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
|
||||
)
|
||||
|
||||
encoder.start_episode(video_keys, tmp_path)
|
||||
@@ -309,7 +303,7 @@ class TestStreamingVideoEncoder:
|
||||
video_keys = [f"{OBS_IMAGES}.cam"]
|
||||
encoder = StreamingVideoEncoder(
|
||||
fps=30,
|
||||
camera_encoder=self._make_encoder_config(
|
||||
rgb_encoder=self._make_encoder_config(
|
||||
vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13
|
||||
),
|
||||
)
|
||||
@@ -346,7 +340,7 @@ class TestStreamingVideoEncoder:
|
||||
video_keys = [f"{OBS_IMAGES}.cam1", f"{OBS_IMAGES}.cam2"]
|
||||
encoder = StreamingVideoEncoder(
|
||||
fps=30,
|
||||
camera_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
|
||||
rgb_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
|
||||
)
|
||||
encoder.start_episode(video_keys, tmp_path)
|
||||
|
||||
@@ -375,7 +369,7 @@ class TestStreamingVideoEncoder:
|
||||
def test_encoder_threads_passed_to_thread(self, tmp_path):
|
||||
"""Test that encoder_threads is stored and passed through to encoder threads."""
|
||||
video_keys = [f"{OBS_IMAGES}.cam"]
|
||||
cfg = VideoEncoderConfig(
|
||||
cfg = RGBEncoderConfig(
|
||||
vcodec="libsvtav1",
|
||||
pix_fmt="yuv420p",
|
||||
g=2,
|
||||
@@ -383,7 +377,7 @@ class TestStreamingVideoEncoder:
|
||||
)
|
||||
encoder = StreamingVideoEncoder(
|
||||
fps=30,
|
||||
camera_encoder=cfg,
|
||||
rgb_encoder=cfg,
|
||||
encoder_threads=2,
|
||||
)
|
||||
assert encoder._encoder_threads == 2
|
||||
@@ -391,7 +385,8 @@ class TestStreamingVideoEncoder:
|
||||
|
||||
# Verify codec options include thread tuning for libsvtav1 (lp=…)
|
||||
thread = encoder._threads[f"{OBS_IMAGES}.cam"]
|
||||
assert "svtav1-params" in thread.codec_options or "threads" in thread.codec_options
|
||||
codec_opts = thread.video_encoder.get_codec_options(encoder_threads=thread.encoder_threads)
|
||||
assert "svtav1-params" in codec_opts or "threads" in codec_opts
|
||||
|
||||
# Feed some frames and finish to ensure it works end-to-end
|
||||
num_frames = 10
|
||||
@@ -422,7 +417,7 @@ class TestStreamingVideoEncoder:
|
||||
video_keys = [f"{OBS_IMAGES}.cam"]
|
||||
encoder = StreamingVideoEncoder(
|
||||
fps=30,
|
||||
camera_encoder=self._make_encoder_config(
|
||||
rgb_encoder=self._make_encoder_config(
|
||||
vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13
|
||||
),
|
||||
queue_maxsize=1,
|
||||
|
||||
@@ -26,7 +26,7 @@ pytest.importorskip("av", reason="av is required (install lerobot[dataset])")
|
||||
|
||||
import av # noqa: E402
|
||||
|
||||
from lerobot.configs import VALID_VIDEO_CODECS, VideoEncoderConfig
|
||||
from lerobot.configs import VALID_VIDEO_CODECS, DepthEncoderConfig, RGBEncoderConfig, VideoEncoderConfig
|
||||
from lerobot.datasets.image_writer import write_image
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
from lerobot.datasets.pyav_utils import get_codec
|
||||
@@ -37,7 +37,15 @@ from lerobot.datasets.video_utils import (
|
||||
get_video_info,
|
||||
reencode_video,
|
||||
)
|
||||
from tests.fixtures.constants import DUMMY_VIDEO_INFO
|
||||
from tests.fixtures.constants import (
|
||||
DUMMY_DEPTH_FEATURES,
|
||||
DUMMY_DEPTH_KEY,
|
||||
DUMMY_DEPTH_VIDEO_INFO_FULL,
|
||||
DUMMY_VIDEO_FEATURES,
|
||||
DUMMY_VIDEO_INFO,
|
||||
DUMMY_VIDEO_KEY,
|
||||
)
|
||||
from tests.fixtures.dataset_factories import add_frames
|
||||
|
||||
|
||||
# Per-codec skip markers — validation tests only fire when the codec is available
|
||||
@@ -48,19 +56,74 @@ def _require_encoder(vcodec: str) -> pytest.MarkDecorator:
|
||||
|
||||
require_libsvtav1 = _require_encoder("libsvtav1")
|
||||
require_h264 = _require_encoder("h264")
|
||||
require_hevc = _require_encoder("hevc")
|
||||
require_videotoolbox = _require_encoder("h264_videotoolbox")
|
||||
require_nvenc = _require_encoder("h264_nvenc")
|
||||
require_vaapi = _require_encoder("h264_vaapi")
|
||||
require_qsv = _require_encoder("h264_qsv")
|
||||
|
||||
|
||||
# ─── VideoEncoderConfig / codec options ──────────────────────────────
|
||||
TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "encoded_videos"
|
||||
|
||||
|
||||
def _write_color_frames(imgs_dir: Path, num_frames: int = 4, height: int = 64, width: int = 96) -> None:
|
||||
imgs_dir.mkdir(parents=True, exist_ok=True)
|
||||
for i in range(num_frames):
|
||||
arr = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
|
||||
write_image(arr, imgs_dir / f"frame-{i:06d}.png")
|
||||
|
||||
|
||||
def _write_depth_frames(imgs_dir: Path, num_frames: int = 4, height: int = 64, width: int = 96) -> None:
|
||||
"""Write synthetic uint16 depth TIFFs (millimetres) for depth encoder tests.
|
||||
|
||||
Uses a smooth linear ramp + per-frame offset (not white noise) so HEVC Main 12
|
||||
on ``gray12le`` compresses well. Values span ~100 mm to 10 m, covering most
|
||||
of the default ``[DEPTH_MIN, DEPTH_MAX]`` metres range after
|
||||
``quantize_depth(input_unit="auto"="mm")``.
|
||||
"""
|
||||
imgs_dir.mkdir(parents=True, exist_ok=True)
|
||||
base = np.linspace(100.0, 10_000.0, height * width, dtype=np.float32).reshape(height, width)
|
||||
for i in range(num_frames):
|
||||
arr = (base + 50.0 * i).clip(0, 65535).astype(np.uint16)
|
||||
write_image(arr, imgs_dir / f"frame-{i:06d}.tiff")
|
||||
|
||||
|
||||
def _encode_video(
|
||||
path: Path,
|
||||
num_frames: int = 4,
|
||||
fps: int = 30,
|
||||
cfg: VideoEncoderConfig | None = None,
|
||||
depth: bool = False,
|
||||
) -> Path:
|
||||
"""Write synthetic frames to a temp dir and encode them to ``path``.
|
||||
|
||||
``depth=False`` writes uint8 RGB PNG noise and encodes with ``cfg``
|
||||
(defaulting to the library default). ``depth=True`` writes synthetic uint16
|
||||
depth TIFFs and encodes with ``cfg`` or a default :class:`DepthEncoderConfig`
|
||||
(HEVC Main 12 / ``gray12le``).
|
||||
"""
|
||||
imgs_dir = path.parent / f"imgs_{path.stem}"
|
||||
if depth:
|
||||
_write_depth_frames(imgs_dir, num_frames=num_frames)
|
||||
cfg = cfg or DepthEncoderConfig()
|
||||
else:
|
||||
_write_color_frames(imgs_dir, num_frames=num_frames)
|
||||
encode_video_frames(imgs_dir, path, fps=fps, video_encoder=cfg, overwrite=True)
|
||||
return path
|
||||
|
||||
|
||||
def _read_feature_info(dataset: LeRobotDataset, key: str = DUMMY_VIDEO_KEY) -> dict:
|
||||
info = json.loads((dataset.root / INFO_PATH).read_text())
|
||||
return info["features"][key]["info"]
|
||||
|
||||
|
||||
# ─── RGBEncoderConfig / codec options ──────────────────────────────
|
||||
|
||||
|
||||
class TestCodecOptions:
|
||||
@require_libsvtav1
|
||||
def test_libsvtav1_defaults(self):
|
||||
cfg = VideoEncoderConfig()
|
||||
cfg = RGBEncoderConfig()
|
||||
opts = cfg.get_codec_options()
|
||||
assert opts["g"] == 2
|
||||
assert opts["crf"] == 30
|
||||
@@ -68,12 +131,12 @@ class TestCodecOptions:
|
||||
|
||||
@require_libsvtav1
|
||||
def test_libsvtav1_custom_preset(self):
|
||||
cfg = VideoEncoderConfig(preset=8)
|
||||
cfg = RGBEncoderConfig(preset=8)
|
||||
assert cfg.get_codec_options()["preset"] == 8
|
||||
|
||||
@require_h264
|
||||
def test_h264_options(self):
|
||||
cfg = VideoEncoderConfig(vcodec="h264", g=10, crf=23, preset=None)
|
||||
cfg = RGBEncoderConfig(vcodec="h264", g=10, crf=23, preset=None)
|
||||
opts = cfg.get_codec_options()
|
||||
assert opts["g"] == 10
|
||||
assert opts["crf"] == 23
|
||||
@@ -81,120 +144,120 @@ class TestCodecOptions:
|
||||
|
||||
@require_videotoolbox
|
||||
def test_videotoolbox_options(self):
|
||||
cfg = VideoEncoderConfig(vcodec="h264_videotoolbox", g=2, crf=30, preset=None)
|
||||
cfg = RGBEncoderConfig(vcodec="h264_videotoolbox", g=2, crf=30, preset=None)
|
||||
opts = cfg.get_codec_options()
|
||||
assert opts["g"] == 2
|
||||
assert opts["q:v"] == 40
|
||||
assert "crf" not in opts
|
||||
|
||||
@_require_encoder("h264_nvenc")
|
||||
@require_nvenc
|
||||
def test_nvenc_options(self):
|
||||
cfg = VideoEncoderConfig(vcodec="h264_nvenc", g=2, crf=25, preset=None)
|
||||
cfg = RGBEncoderConfig(vcodec="h264_nvenc", g=2, crf=25, preset=None)
|
||||
opts = cfg.get_codec_options()
|
||||
assert opts["rc"] == 0
|
||||
assert opts["qp"] == 25
|
||||
assert "crf" not in opts
|
||||
assert opts["g"] == 2
|
||||
|
||||
@_require_encoder("h264_vaapi")
|
||||
@require_vaapi
|
||||
def test_vaapi_options(self):
|
||||
cfg = VideoEncoderConfig(vcodec="h264_vaapi", crf=28, preset=None)
|
||||
cfg = RGBEncoderConfig(vcodec="h264_vaapi", crf=28, preset=None)
|
||||
assert cfg.get_codec_options()["qp"] == 28
|
||||
|
||||
@_require_encoder("h264_qsv")
|
||||
@require_qsv
|
||||
def test_qsv_options(self):
|
||||
cfg = VideoEncoderConfig(vcodec="h264_qsv", crf=25, preset=None)
|
||||
cfg = RGBEncoderConfig(vcodec="h264_qsv", crf=25, preset=None)
|
||||
assert cfg.get_codec_options()["global_quality"] == 25
|
||||
|
||||
@require_h264
|
||||
def test_no_g_no_crf(self):
|
||||
cfg = VideoEncoderConfig(vcodec="h264", g=None, crf=None, preset=None)
|
||||
cfg = RGBEncoderConfig(vcodec="h264", g=None, crf=None, preset=None)
|
||||
opts = cfg.get_codec_options()
|
||||
assert "g" not in opts
|
||||
assert "crf" not in opts
|
||||
|
||||
@require_libsvtav1
|
||||
def test_encoder_threads_libsvtav1(self):
|
||||
cfg = VideoEncoderConfig(fast_decode=0)
|
||||
cfg = RGBEncoderConfig(fast_decode=0)
|
||||
opts = cfg.get_codec_options(encoder_threads=4)
|
||||
assert "lp=4" in opts.get("svtav1-params", "")
|
||||
|
||||
@require_h264
|
||||
def test_encoder_threads_h264(self):
|
||||
cfg = VideoEncoderConfig(vcodec="h264", preset=None)
|
||||
cfg = RGBEncoderConfig(vcodec="h264", preset=None)
|
||||
assert cfg.get_codec_options(encoder_threads=2)["threads"] == 2
|
||||
|
||||
@require_libsvtav1
|
||||
def test_fast_decode_libsvtav1(self):
|
||||
cfg = VideoEncoderConfig(fast_decode=1)
|
||||
cfg = RGBEncoderConfig(fast_decode=1)
|
||||
opts = cfg.get_codec_options()
|
||||
assert "fast-decode=1" in opts.get("svtav1-params", "")
|
||||
|
||||
@require_libsvtav1
|
||||
def test_libsvtav1_fast_decode_clamped_to_svt_range(self):
|
||||
"""Out-of-range fast_decode is clamped to [0, 2] in svtav1-params (SVT-AV1 FastDecode)."""
|
||||
cfg = VideoEncoderConfig(fast_decode=100)
|
||||
cfg = RGBEncoderConfig(fast_decode=100)
|
||||
assert "fast-decode=2" in cfg.get_codec_options().get("svtav1-params", "")
|
||||
cfg_neg = VideoEncoderConfig(fast_decode=-5)
|
||||
cfg_neg = RGBEncoderConfig(fast_decode=-5)
|
||||
assert "fast-decode=0" in cfg_neg.get_codec_options().get("svtav1-params", "")
|
||||
|
||||
@require_h264
|
||||
def test_fast_decode_h264(self):
|
||||
cfg = VideoEncoderConfig(vcodec="h264", fast_decode=1, preset=None)
|
||||
cfg = RGBEncoderConfig(vcodec="h264", fast_decode=1, preset=None)
|
||||
assert cfg.get_codec_options()["tune"] == "fastdecode"
|
||||
|
||||
@require_libsvtav1
|
||||
def test_pix_fmt_unsupported_raises(self):
|
||||
"""Passing an unsupported pix_fmt is a hard error."""
|
||||
with pytest.raises(ValueError, match="pix_fmt"):
|
||||
VideoEncoderConfig(pix_fmt="yuv444p") # libsvtav1 only supports yuv420p variants
|
||||
RGBEncoderConfig(pix_fmt="yuv444p") # libsvtav1 only supports yuv420p variants
|
||||
|
||||
@require_libsvtav1
|
||||
@require_h264
|
||||
def test_preset_default_behaviour(self):
|
||||
"""Empty constructor picks preset=12 (libsvtav1 path); other codecs stay None."""
|
||||
assert VideoEncoderConfig().preset == 12
|
||||
assert VideoEncoderConfig(vcodec="libsvtav1").preset == 12
|
||||
assert VideoEncoderConfig(vcodec="h264").preset is None
|
||||
assert VideoEncoderConfig(vcodec="h264", preset=None).preset is None
|
||||
assert RGBEncoderConfig().preset == 12
|
||||
assert RGBEncoderConfig(vcodec="libsvtav1").preset == 12
|
||||
assert RGBEncoderConfig(vcodec="h264").preset is None
|
||||
assert RGBEncoderConfig(vcodec="h264", preset=None).preset is None
|
||||
|
||||
@require_h264
|
||||
def test_preset_string_on_h264(self):
|
||||
"""h264 accepts string presets and forwards them to FFmpeg."""
|
||||
cfg = VideoEncoderConfig(vcodec="h264", preset="slow")
|
||||
cfg = RGBEncoderConfig(vcodec="h264", preset="slow")
|
||||
assert cfg.get_codec_options()["preset"] == "slow"
|
||||
|
||||
@require_videotoolbox
|
||||
def test_preset_on_videotoolbox_not_set(self):
|
||||
"""videotoolbox has no preset option at all."""
|
||||
cfg = VideoEncoderConfig(vcodec="h264_videotoolbox", preset="slow")
|
||||
cfg = RGBEncoderConfig(vcodec="h264_videotoolbox", preset="slow")
|
||||
assert "preset" not in cfg.get_codec_options()
|
||||
|
||||
@require_libsvtav1
|
||||
def test_libsvtav1_preset_out_of_range_raises(self):
|
||||
"""libsvtav1 preset must sit in [-2, 13] as exposed by PyAV."""
|
||||
with pytest.raises(ValueError, match="out of range"):
|
||||
VideoEncoderConfig(vcodec="libsvtav1", preset=100)
|
||||
RGBEncoderConfig(vcodec="libsvtav1", preset=100)
|
||||
with pytest.raises(ValueError, match="out of range"):
|
||||
VideoEncoderConfig(vcodec="libsvtav1", preset=-3)
|
||||
RGBEncoderConfig(vcodec="libsvtav1", preset=-3)
|
||||
|
||||
@require_libsvtav1
|
||||
def test_libsvtav1_crf_out_of_range_raises(self):
|
||||
"""libsvtav1 crf must sit in [0, 63]."""
|
||||
with pytest.raises(ValueError, match="crf.*out of range"):
|
||||
VideoEncoderConfig(vcodec="libsvtav1", crf=64)
|
||||
RGBEncoderConfig(vcodec="libsvtav1", crf=64)
|
||||
|
||||
@require_libsvtav1
|
||||
def test_libsvtav1_crf_rejects_python_float(self):
|
||||
"""libsvtav1 exposes ``crf`` as an INT AVOption; Python float must not pass validation."""
|
||||
with pytest.raises(ValueError, match="float values are not allowed"):
|
||||
VideoEncoderConfig(vcodec="libsvtav1", crf=2.5)
|
||||
RGBEncoderConfig(vcodec="libsvtav1", crf=2.5)
|
||||
|
||||
@require_libsvtav1
|
||||
def test_libsvtav1_extra_crf_rejects_fractional_string(self):
|
||||
"""INT options reject fractional values even when supplied only via ``extra_options``."""
|
||||
with pytest.raises(ValueError, match="float values are not allowed"):
|
||||
VideoEncoderConfig(
|
||||
RGBEncoderConfig(
|
||||
vcodec="libsvtav1",
|
||||
crf=None,
|
||||
extra_options={"crf": "2.5"},
|
||||
@@ -203,7 +266,7 @@ class TestCodecOptions:
|
||||
@require_libsvtav1
|
||||
def test_libsvtav1_extra_crf_rejects_float(self):
|
||||
with pytest.raises(ValueError, match="float values are not allowed"):
|
||||
VideoEncoderConfig(
|
||||
RGBEncoderConfig(
|
||||
vcodec="libsvtav1",
|
||||
crf=None,
|
||||
extra_options={"crf": 2.5},
|
||||
@@ -212,13 +275,13 @@ class TestCodecOptions:
|
||||
@require_h264
|
||||
def test_h264_crf_accepts_float_and_int(self):
|
||||
"""x264 exposes crf as a FLOAT option, so both int and float are accepted."""
|
||||
assert VideoEncoderConfig(vcodec="h264", crf=23).get_codec_options()["crf"] == 23
|
||||
assert VideoEncoderConfig(vcodec="h264", crf=23.5).get_codec_options()["crf"] == 23.5
|
||||
assert RGBEncoderConfig(vcodec="h264", crf=23).get_codec_options()["crf"] == 23
|
||||
assert RGBEncoderConfig(vcodec="h264", crf=23.5).get_codec_options()["crf"] == 23.5
|
||||
|
||||
@require_libsvtav1
|
||||
def test_validate_is_rerunnable(self):
|
||||
"""After mutating a field, validate() re-checks and surfaces new issues."""
|
||||
cfg = VideoEncoderConfig(vcodec="libsvtav1")
|
||||
cfg = RGBEncoderConfig(vcodec="libsvtav1")
|
||||
cfg.preset = 100 # now out of range
|
||||
with pytest.raises(ValueError, match="out of range"):
|
||||
cfg.validate()
|
||||
@@ -227,58 +290,58 @@ class TestCodecOptions:
|
||||
class TestExtraOptions:
|
||||
@require_libsvtav1
|
||||
def test_default_is_empty_dict(self):
|
||||
cfg = VideoEncoderConfig()
|
||||
cfg = RGBEncoderConfig()
|
||||
assert cfg.extra_options == {}
|
||||
|
||||
@require_libsvtav1
|
||||
def test_unknown_key_passes_through(self):
|
||||
"""Keys not published as AVOptions are forwarded to FFmpeg."""
|
||||
cfg = VideoEncoderConfig(extra_options={"totally_made_up_option": "value"})
|
||||
cfg = RGBEncoderConfig(extra_options={"totally_made_up_option": "value"})
|
||||
assert cfg.extra_options == {"totally_made_up_option": "value"}
|
||||
|
||||
@require_libsvtav1
|
||||
def test_numeric_value_in_range_ok(self):
|
||||
"""libsvtav1 exposes ``qp`` as INT in [0, 63]."""
|
||||
cfg = VideoEncoderConfig(extra_options={"qp": 30})
|
||||
cfg = RGBEncoderConfig(extra_options={"qp": 30})
|
||||
assert cfg.extra_options == {"qp": 30}
|
||||
|
||||
@require_libsvtav1
|
||||
def test_numeric_out_of_range_raises(self):
|
||||
with pytest.raises(ValueError, match=r"qp=.*out of range"):
|
||||
VideoEncoderConfig(extra_options={"qp": 999})
|
||||
RGBEncoderConfig(extra_options={"qp": 999})
|
||||
|
||||
@require_libsvtav1
|
||||
def test_numeric_string_accepted_in_range(self):
|
||||
"""Numeric strings are accepted for numeric options (mirrors FFmpeg)."""
|
||||
cfg = VideoEncoderConfig(extra_options={"qp": "18"})
|
||||
cfg = RGBEncoderConfig(extra_options={"qp": "18"})
|
||||
assert cfg.extra_options == {"qp": "18"}
|
||||
|
||||
@require_libsvtav1
|
||||
def test_numeric_string_out_of_range_raises(self):
|
||||
with pytest.raises(ValueError, match=r"qp=.*out of range"):
|
||||
VideoEncoderConfig(extra_options={"qp": "999"})
|
||||
RGBEncoderConfig(extra_options={"qp": "999"})
|
||||
|
||||
@require_libsvtav1
|
||||
def test_non_numeric_string_on_numeric_option_raises(self):
|
||||
with pytest.raises(ValueError, match=r"qp=.*not numeric"):
|
||||
VideoEncoderConfig(extra_options={"qp": "medium"})
|
||||
RGBEncoderConfig(extra_options={"qp": "medium"})
|
||||
|
||||
@require_libsvtav1
|
||||
def test_bool_on_numeric_option_raises(self):
|
||||
"""``bool`` is explicitly rejected for numeric options."""
|
||||
with pytest.raises(ValueError, match=r"qp=.*not numeric"):
|
||||
VideoEncoderConfig(extra_options={"qp": True})
|
||||
RGBEncoderConfig(extra_options={"qp": True})
|
||||
|
||||
@require_h264
|
||||
def test_string_option_passes_through_unchecked(self):
|
||||
"""String-typed AVOptions are NOT enum-checked (too many accept freeform)."""
|
||||
cfg = VideoEncoderConfig(vcodec="h264", preset=None, extra_options={"tune": "some-future-tune"})
|
||||
cfg = RGBEncoderConfig(vcodec="h264", preset=None, extra_options={"tune": "some-future-tune"})
|
||||
assert cfg.extra_options == {"tune": "some-future-tune"}
|
||||
|
||||
@require_libsvtav1
|
||||
def test_merged_into_codec_options_and_stringified(self):
|
||||
"""Typed merge by default; ``as_strings=True`` matches FFmpeg option dict."""
|
||||
cfg = VideoEncoderConfig(extra_options={"qp": 20})
|
||||
cfg = RGBEncoderConfig(extra_options={"qp": 20})
|
||||
opts = cfg.get_codec_options()
|
||||
assert opts["qp"] == 20
|
||||
assert isinstance(opts["qp"], int)
|
||||
@@ -287,25 +350,25 @@ class TestExtraOptions:
|
||||
@require_libsvtav1
|
||||
def test_structured_fields_win_on_collision(self):
|
||||
"""A colliding extra_options key is discarded; the structured field wins."""
|
||||
cfg = VideoEncoderConfig(crf=30, extra_options={"crf": 18})
|
||||
cfg = RGBEncoderConfig(crf=30, extra_options={"crf": 18})
|
||||
assert cfg.get_codec_options()["crf"] == 30
|
||||
|
||||
|
||||
class TestEncoderDetection:
|
||||
@require_h264
|
||||
def test_explicit_codec_kept_when_available(self):
|
||||
cfg = VideoEncoderConfig(vcodec="h264")
|
||||
cfg = RGBEncoderConfig(vcodec="h264")
|
||||
assert cfg.vcodec == "h264"
|
||||
|
||||
@require_videotoolbox
|
||||
def test_auto_picks_videotoolbox_when_available(self):
|
||||
"""``h264_videotoolbox`` sits at the top of ``HW_VIDEO_CODECS`` so it wins when present."""
|
||||
cfg = VideoEncoderConfig(vcodec="auto")
|
||||
cfg = RGBEncoderConfig(vcodec="auto")
|
||||
assert cfg.vcodec == "h264_videotoolbox"
|
||||
|
||||
def test_invalid_codec_raises(self):
|
||||
with pytest.raises(ValueError, match="Invalid vcodec"):
|
||||
VideoEncoderConfig(vcodec="not_a_real_codec")
|
||||
RGBEncoderConfig(vcodec="not_a_real_codec")
|
||||
|
||||
def test_hw_encoder_names_listed_as_valid(self):
|
||||
assert "auto" in VALID_VIDEO_CODECS
|
||||
@@ -313,59 +376,6 @@ class TestEncoderDetection:
|
||||
assert "h264_nvenc" in VALID_VIDEO_CODECS
|
||||
|
||||
|
||||
TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "encoded_videos"
|
||||
|
||||
# Default video feature set used by persistence tests.
|
||||
VIDEO_FEATURES = {
|
||||
"observation.images.cam": {
|
||||
"dtype": "video",
|
||||
"shape": (64, 96, 3),
|
||||
"names": ["height", "width", "channels"],
|
||||
},
|
||||
"action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]},
|
||||
}
|
||||
VIDEO_KEY = "observation.images.cam"
|
||||
|
||||
|
||||
def _write_frames(imgs_dir: Path, num_frames: int = 4, height: int = 64, width: int = 96) -> None:
|
||||
imgs_dir.mkdir(parents=True, exist_ok=True)
|
||||
for i in range(num_frames):
|
||||
arr = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
|
||||
write_image(arr, imgs_dir / f"frame-{i:06d}.png")
|
||||
|
||||
|
||||
def _encode_video(
|
||||
path: Path, num_frames: int = 4, fps: int = 30, cfg: VideoEncoderConfig | None = None
|
||||
) -> Path:
|
||||
imgs_dir = path.parent / f"imgs_{path.stem}"
|
||||
_write_frames(imgs_dir, num_frames=num_frames)
|
||||
encode_video_frames(imgs_dir, path, fps=fps, camera_encoder=cfg, overwrite=True)
|
||||
return path
|
||||
|
||||
|
||||
def _read_feature_info(dataset: LeRobotDataset) -> dict:
|
||||
info = json.loads((dataset.root / INFO_PATH).read_text())
|
||||
return info["features"][VIDEO_KEY]["info"]
|
||||
|
||||
|
||||
def _add_frames(dataset: LeRobotDataset, num_frames: int, video_keys: list[str] | None = None) -> None:
|
||||
from lerobot.utils.constants import DEFAULT_FEATURES
|
||||
|
||||
if video_keys is None:
|
||||
video_keys = dataset.meta.video_keys
|
||||
for _ in range(num_frames):
|
||||
frame: dict = {"task": "test"}
|
||||
for key, ft in dataset.meta.features.items():
|
||||
if key in DEFAULT_FEATURES:
|
||||
continue
|
||||
shape = ft["shape"]
|
||||
if key in video_keys:
|
||||
frame[key] = np.random.randint(0, 256, shape, dtype=np.uint8)
|
||||
else:
|
||||
frame[key] = np.zeros(shape, dtype=np.float32)
|
||||
dataset.add_frame(frame)
|
||||
|
||||
|
||||
class TestGetVideoInfo:
|
||||
def test_returns_all_stream_fields(self):
|
||||
info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4")
|
||||
@@ -375,7 +385,7 @@ class TestGetVideoInfo:
|
||||
assert info["video.pix_fmt"] == "yuv420p"
|
||||
assert info["video.fps"] == 30
|
||||
assert info["video.channels"] == 3
|
||||
assert info["video.is_depth_map"] is False
|
||||
assert info["is_depth_map"] is False
|
||||
assert info["has_audio"] is False
|
||||
assert "video.g" not in info
|
||||
assert "video.crf" not in info
|
||||
@@ -383,9 +393,9 @@ class TestGetVideoInfo:
|
||||
|
||||
@require_libsvtav1
|
||||
def test_merges_encoder_config_as_video_prefixed_entries(self):
|
||||
cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
|
||||
cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
|
||||
|
||||
info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", camera_encoder=cfg)
|
||||
info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", video_encoder=cfg)
|
||||
|
||||
assert info["video.g"] == 2
|
||||
assert info["video.crf"] == 30
|
||||
@@ -396,13 +406,18 @@ class TestGetVideoInfo:
|
||||
|
||||
@require_libsvtav1
|
||||
def test_stream_derived_keys_take_precedence_over_config(self):
|
||||
cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p")
|
||||
cfg = RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p")
|
||||
|
||||
info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", camera_encoder=cfg)
|
||||
info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", video_encoder=cfg)
|
||||
|
||||
assert info["video.codec"] # populated from stream, not from config's vcodec
|
||||
assert info["video.pix_fmt"] == "yuv420p"
|
||||
|
||||
def test_depth_encoder_config_sets_is_depth_map_true(self):
|
||||
"""A ``DepthEncoderConfig`` causes ``get_video_info`` to mark the stream as depth."""
|
||||
info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", video_encoder=DepthEncoderConfig())
|
||||
assert info["is_depth_map"] is True
|
||||
|
||||
|
||||
class TestEncodeVideoFrames:
|
||||
@require_libsvtav1
|
||||
@@ -434,7 +449,7 @@ class TestEncodeVideoFrames:
|
||||
|
||||
def test_overwrite_false_skips_existing_file(self, tmp_path):
|
||||
imgs_dir = tmp_path / "imgs"
|
||||
_write_frames(imgs_dir)
|
||||
_write_color_frames(imgs_dir)
|
||||
video_path = tmp_path / "out.mp4"
|
||||
sentinel = b"pre-existing content"
|
||||
video_path.write_bytes(sentinel)
|
||||
@@ -446,7 +461,7 @@ class TestEncodeVideoFrames:
|
||||
@require_libsvtav1
|
||||
def test_overwrite_true_replaces_existing_file(self, tmp_path):
|
||||
imgs_dir = tmp_path / "imgs"
|
||||
_write_frames(imgs_dir)
|
||||
_write_color_frames(imgs_dir)
|
||||
video_path = tmp_path / "out.mp4"
|
||||
video_path.write_bytes(b"stale content")
|
||||
|
||||
@@ -458,10 +473,10 @@ class TestEncodeVideoFrames:
|
||||
@require_libsvtav1
|
||||
def test_custom_encoder_config_fields_stored_in_info(self, tmp_path):
|
||||
"""All stream-derived and encoder config fields are present after encoding."""
|
||||
cfg = VideoEncoderConfig(vcodec="libsvtav1", g=4, crf=25, preset=10)
|
||||
cfg = RGBEncoderConfig(vcodec="libsvtav1", g=4, crf=25, preset=10)
|
||||
video_path = _encode_video(tmp_path / "out.mp4", num_frames=4, fps=30, cfg=cfg)
|
||||
|
||||
info = get_video_info(video_path, camera_encoder=cfg)
|
||||
info = get_video_info(video_path, video_encoder=cfg)
|
||||
|
||||
# Stream-derived
|
||||
assert info["video.height"] == 64
|
||||
@@ -470,7 +485,7 @@ class TestEncodeVideoFrames:
|
||||
assert info["video.codec"] == "av1"
|
||||
assert info["video.pix_fmt"] == "yuv420p"
|
||||
assert info["video.fps"] == 30
|
||||
assert info["video.is_depth_map"] is False
|
||||
assert info["is_depth_map"] is False
|
||||
assert info["has_audio"] is False
|
||||
# Encoder config
|
||||
assert info["video.g"] == 4
|
||||
@@ -487,15 +502,15 @@ class TestReencodeVideo:
|
||||
def test_reencode_video(self, tmp_path):
|
||||
src = TEST_ARTIFACTS_DIR / "clip_4frames.mp4"
|
||||
out = tmp_path / "reencoded.mp4"
|
||||
cfg = VideoEncoderConfig(vcodec="h264", g=6, crf=23, pix_fmt="yuv444p")
|
||||
reencode_video(src, out, camera_encoder=cfg, overwrite=True)
|
||||
cfg = RGBEncoderConfig(vcodec="h264", g=6, crf=23, pix_fmt="yuv444p")
|
||||
reencode_video(src, out, video_encoder=cfg, overwrite=True)
|
||||
|
||||
assert out.exists()
|
||||
with av.open(str(out)) as container:
|
||||
n_frames = sum(1 for _ in container.decode(video=0))
|
||||
assert n_frames == 4
|
||||
|
||||
info = get_video_info(out, camera_encoder=cfg)
|
||||
info = get_video_info(out, video_encoder=cfg)
|
||||
assert info["video.codec"] == "h264"
|
||||
assert info["video.pix_fmt"] == "yuv444p"
|
||||
assert info["video.height"] == 64
|
||||
@@ -508,8 +523,8 @@ class TestReencodeVideo:
|
||||
def test_reencode_video_trim_window(self, tmp_path):
|
||||
src = TEST_ARTIFACTS_DIR / "clip_6frames.mp4"
|
||||
out = tmp_path / "trim_window.mp4"
|
||||
cfg = VideoEncoderConfig(vcodec="h264")
|
||||
reencode_video(src, out, camera_encoder=cfg, start_time_s=0.05, end_time_s=0.12, overwrite=True)
|
||||
cfg = RGBEncoderConfig(vcodec="h264")
|
||||
reencode_video(src, out, video_encoder=cfg, start_time_s=0.05, end_time_s=0.12, overwrite=True)
|
||||
|
||||
with av.open(str(out)) as container:
|
||||
frames = list(container.decode(video=0))
|
||||
@@ -578,12 +593,12 @@ class TestEncoderConfigPersistence:
|
||||
|
||||
@require_libsvtav1
|
||||
def test_first_episode_save_persists_encoder_config(self, tmp_path, empty_lerobot_dataset_factory):
|
||||
cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
|
||||
cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
|
||||
dataset = empty_lerobot_dataset_factory(
|
||||
root=tmp_path / "ds", features=VIDEO_FEATURES, use_videos=True, camera_encoder=cfg
|
||||
root=tmp_path / "ds", features=DUMMY_VIDEO_FEATURES, use_videos=True, rgb_encoder=cfg
|
||||
)
|
||||
|
||||
_add_frames(dataset, num_frames=4)
|
||||
add_frames(dataset, num_frames=4)
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
|
||||
@@ -601,16 +616,16 @@ class TestEncoderConfigPersistence:
|
||||
|
||||
@require_libsvtav1
|
||||
def test_second_episode_does_not_overwrite_encoder_fields(self, tmp_path, empty_lerobot_dataset_factory):
|
||||
cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
|
||||
cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
|
||||
dataset = empty_lerobot_dataset_factory(
|
||||
root=tmp_path / "ds", features=VIDEO_FEATURES, use_videos=True, camera_encoder=cfg
|
||||
root=tmp_path / "ds", features=DUMMY_VIDEO_FEATURES, use_videos=True, rgb_encoder=cfg
|
||||
)
|
||||
|
||||
_add_frames(dataset, num_frames=4)
|
||||
add_frames(dataset, num_frames=4)
|
||||
dataset.save_episode()
|
||||
first_info = dict(_read_feature_info(dataset))
|
||||
|
||||
_add_frames(dataset, num_frames=4)
|
||||
add_frames(dataset, num_frames=4)
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
|
||||
@@ -618,13 +633,13 @@ class TestEncoderConfigPersistence:
|
||||
|
||||
|
||||
class TestFromVideoInfo:
|
||||
"""``VideoEncoderConfig.from_video_info`` reconstructs an encoder config
|
||||
"""``RGBEncoderConfig.from_video_info`` reconstructs an encoder config
|
||||
from the ``video.*`` keys persisted in a dataset's ``info.json``.
|
||||
"""
|
||||
|
||||
@require_libsvtav1
|
||||
def test_reconstructs_from_dummy_video_info(self):
|
||||
cfg = VideoEncoderConfig.from_video_info(DUMMY_VIDEO_INFO)
|
||||
cfg = RGBEncoderConfig.from_video_info(DUMMY_VIDEO_INFO)
|
||||
|
||||
# Canonical stream codec ``"av1"`` is aliased to the encoder name.
|
||||
assert cfg.vcodec == "libsvtav1"
|
||||
@@ -636,4 +651,220 @@ class TestFromVideoInfo:
|
||||
assert cfg.video_backend == DUMMY_VIDEO_INFO["video.video_backend"]
|
||||
# ``{}`` placeholder (typical after a merge with disagreeing sources)
|
||||
# must not leak into the reconstructed config.
|
||||
assert cfg.extra_options == VideoEncoderConfig().extra_options
|
||||
assert cfg.extra_options == RGBEncoderConfig().extra_options
|
||||
|
||||
|
||||
# ─── Depth-specific encoding tests ────────────────────────────────────
|
||||
|
||||
|
||||
class TestEncodeDepthVideoFrames:
|
||||
"""Depth mirror of :class:`TestEncodeVideoFrames`.
|
||||
|
||||
Exercises ``encode_video_frames`` end-to-end through
|
||||
:class:`DepthEncoderConfig` (HEVC Main 12 / ``gray12le``) on synthetic
|
||||
uint16 depth TIFFs.
|
||||
"""
|
||||
|
||||
@require_hevc
|
||||
def test_produces_readable_file(self, tmp_path):
|
||||
video_path = _encode_video(tmp_path / "out.mp4", depth=True)
|
||||
|
||||
assert video_path.exists()
|
||||
info = get_video_info(video_path, video_encoder=DepthEncoderConfig())
|
||||
assert info["video.height"] == 64
|
||||
assert info["video.width"] == 96
|
||||
assert info["video.codec"] == "hevc"
|
||||
assert info["video.pix_fmt"] == "gray12le"
|
||||
assert info["video.channels"] == 1
|
||||
assert info["is_depth_map"] is True
|
||||
|
||||
@require_hevc
|
||||
def test_frame_count_and_duration_match_input(self, tmp_path):
|
||||
num_frames = 10
|
||||
fps = 30
|
||||
video_path = _encode_video(tmp_path / "out.mp4", num_frames=num_frames, fps=fps, depth=True)
|
||||
|
||||
with av.open(str(video_path)) as container:
|
||||
stream = container.streams.video[0]
|
||||
actual_frames = sum(1 for _ in container.decode(stream))
|
||||
duration = (
|
||||
float(stream.duration * stream.time_base)
|
||||
if stream.duration is not None
|
||||
else float(container.duration / av.time_base)
|
||||
)
|
||||
|
||||
assert actual_frames == num_frames
|
||||
assert abs(duration - num_frames / fps) < 0.1
|
||||
|
||||
def test_overwrite_false_skips_existing_file(self, tmp_path):
|
||||
"""Codec-agnostic: file-system semantics must hold even without an HEVC encoder."""
|
||||
imgs_dir = tmp_path / "imgs"
|
||||
_write_depth_frames(imgs_dir)
|
||||
video_path = tmp_path / "out.mp4"
|
||||
sentinel = b"pre-existing depth content"
|
||||
video_path.write_bytes(sentinel)
|
||||
|
||||
encode_video_frames(imgs_dir, video_path, fps=30, video_encoder=DepthEncoderConfig(), overwrite=False)
|
||||
|
||||
assert video_path.read_bytes() == sentinel
|
||||
|
||||
@require_hevc
|
||||
def test_overwrite_true_replaces_existing_file(self, tmp_path):
|
||||
imgs_dir = tmp_path / "imgs"
|
||||
_write_depth_frames(imgs_dir)
|
||||
video_path = tmp_path / "out.mp4"
|
||||
video_path.write_bytes(b"stale content")
|
||||
|
||||
encode_video_frames(imgs_dir, video_path, fps=30, video_encoder=DepthEncoderConfig(), overwrite=True)
|
||||
|
||||
info = get_video_info(video_path, video_encoder=DepthEncoderConfig())
|
||||
assert info["video.height"] == 64
|
||||
assert info["video.pix_fmt"] == "gray12le"
|
||||
assert info["is_depth_map"] is True
|
||||
|
||||
@require_hevc
|
||||
def test_custom_encoder_config_fields_stored_in_info(self, tmp_path):
|
||||
"""All stream-derived and depth-encoder config fields are present after encoding."""
|
||||
cfg = DepthEncoderConfig(
|
||||
vcodec="hevc",
|
||||
pix_fmt="gray12le",
|
||||
g=4,
|
||||
crf=25,
|
||||
extra_options={},
|
||||
depth_min=0.05,
|
||||
depth_max=8.0,
|
||||
shift=2.5,
|
||||
use_log=False,
|
||||
)
|
||||
video_path = _encode_video(tmp_path / "out.mp4", num_frames=4, fps=30, cfg=cfg, depth=True)
|
||||
|
||||
info = get_video_info(video_path, video_encoder=cfg)
|
||||
|
||||
# Stream-derived
|
||||
assert info["video.height"] == 64
|
||||
assert info["video.width"] == 96
|
||||
assert info["video.channels"] == 1
|
||||
assert info["video.codec"] == "hevc"
|
||||
assert info["video.pix_fmt"] == "gray12le"
|
||||
assert info["video.fps"] == 30
|
||||
assert info["is_depth_map"] is True
|
||||
assert info["has_audio"] is False
|
||||
# Base encoder config
|
||||
assert info["video.g"] == 4
|
||||
assert info["video.crf"] == 25
|
||||
assert info["video.fast_decode"] == 0
|
||||
assert info["video.video_backend"] == "pyav"
|
||||
assert info["video.extra_options"] == {}
|
||||
# Depth-specific tuning
|
||||
assert info["video.depth_min"] == 0.05
|
||||
assert info["video.depth_max"] == 8.0
|
||||
assert info["video.shift"] == 2.5
|
||||
assert info["video.use_log"] is False
|
||||
|
||||
|
||||
class TestDepthEncoderConfigPersistence:
|
||||
"""Depth mirror of :class:`TestEncoderConfigPersistence`.
|
||||
|
||||
``DepthEncoderConfig`` must be stored as ``video.<field>`` entries
|
||||
(including the depth-specific ``depth_min`` / ``depth_max`` / ``shift`` /
|
||||
``use_log``) under ``info["features"][<depth_key>]["info"]`` when the
|
||||
first episode is saved.
|
||||
"""
|
||||
|
||||
@require_hevc
|
||||
def test_first_episode_save_persists_depth_encoder_config(self, tmp_path, empty_lerobot_dataset_factory):
|
||||
cfg = DepthEncoderConfig(
|
||||
vcodec="hevc",
|
||||
pix_fmt="gray12le",
|
||||
g=2,
|
||||
crf=30,
|
||||
extra_options={},
|
||||
depth_min=0.05,
|
||||
depth_max=8.0,
|
||||
shift=2.5,
|
||||
use_log=False,
|
||||
)
|
||||
dataset = empty_lerobot_dataset_factory(
|
||||
root=tmp_path / "ds", features=DUMMY_DEPTH_FEATURES, use_videos=True, depth_encoder=cfg
|
||||
)
|
||||
|
||||
add_frames(dataset, num_frames=4)
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
|
||||
info = _read_feature_info(dataset, key=DUMMY_DEPTH_KEY)
|
||||
|
||||
# Stream-derived
|
||||
assert info["video.height"] == 64
|
||||
assert info["video.width"] == 96
|
||||
assert info["video.fps"] == 30
|
||||
assert info["video.codec"] == "hevc"
|
||||
assert info["video.pix_fmt"] == "gray12le"
|
||||
assert info["is_depth_map"] is True
|
||||
# Base encoder config
|
||||
assert info["video.g"] == 2
|
||||
assert info["video.crf"] == 30
|
||||
assert info["video.fast_decode"] == 0
|
||||
assert info["video.video_backend"] == "pyav"
|
||||
assert info["video.extra_options"] == {}
|
||||
# Depth-specific tuning
|
||||
assert info["video.depth_min"] == 0.05
|
||||
assert info["video.depth_max"] == 8.0
|
||||
assert info["video.shift"] == 2.5
|
||||
assert info["video.use_log"] is False
|
||||
|
||||
@require_hevc
|
||||
def test_second_episode_does_not_overwrite_depth_encoder_fields(
|
||||
self, tmp_path, empty_lerobot_dataset_factory
|
||||
):
|
||||
cfg = DepthEncoderConfig(
|
||||
vcodec="hevc",
|
||||
pix_fmt="gray12le",
|
||||
g=2,
|
||||
crf=30,
|
||||
depth_min=0.05,
|
||||
depth_max=8.0,
|
||||
shift=2.5,
|
||||
use_log=False,
|
||||
)
|
||||
dataset = empty_lerobot_dataset_factory(
|
||||
root=tmp_path / "ds", features=DUMMY_DEPTH_FEATURES, use_videos=True, depth_encoder=cfg
|
||||
)
|
||||
|
||||
add_frames(dataset, num_frames=4)
|
||||
dataset.save_episode()
|
||||
first_info = dict(_read_feature_info(dataset, key=DUMMY_DEPTH_KEY))
|
||||
|
||||
add_frames(dataset, num_frames=4)
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
|
||||
assert _read_feature_info(dataset, key=DUMMY_DEPTH_KEY) == first_info
|
||||
|
||||
|
||||
class TestDepthFromVideoInfo:
|
||||
"""``DepthEncoderConfig.from_video_info`` reconstructs a depth encoder
|
||||
config from the ``video.*`` keys persisted in a dataset's ``info.json``.
|
||||
|
||||
Depth mirror of :class:`TestFromVideoInfo`.
|
||||
"""
|
||||
|
||||
@require_hevc
|
||||
def test_reconstructs_from_dummy_depth_video_info(self):
|
||||
cfg = DepthEncoderConfig.from_video_info(DUMMY_DEPTH_VIDEO_INFO_FULL)
|
||||
|
||||
# No alias for ``"hevc"``; the canonical stream codec is reused as-is.
|
||||
assert cfg.vcodec == "hevc"
|
||||
assert cfg.pix_fmt == DUMMY_DEPTH_VIDEO_INFO_FULL["video.pix_fmt"]
|
||||
assert cfg.g == DUMMY_DEPTH_VIDEO_INFO_FULL["video.g"]
|
||||
assert cfg.crf == DUMMY_DEPTH_VIDEO_INFO_FULL["video.crf"]
|
||||
assert cfg.fast_decode == DUMMY_DEPTH_VIDEO_INFO_FULL["video.fast_decode"]
|
||||
assert cfg.video_backend == DUMMY_DEPTH_VIDEO_INFO_FULL["video.video_backend"]
|
||||
# ``{}`` placeholder (typical after a merge with disagreeing sources)
|
||||
# must not leak into the reconstructed config.
|
||||
assert cfg.extra_options == DepthEncoderConfig().extra_options
|
||||
# Depth-specific tuning round-trips through ``info.json``.
|
||||
assert cfg.depth_min == DUMMY_DEPTH_VIDEO_INFO_FULL["video.depth_min"]
|
||||
assert cfg.depth_max == DUMMY_DEPTH_VIDEO_INFO_FULL["video.depth_max"]
|
||||
assert cfg.shift == DUMMY_DEPTH_VIDEO_INFO_FULL["video.shift"]
|
||||
assert cfg.use_log == DUMMY_DEPTH_VIDEO_INFO_FULL["video.use_log"]
|
||||
|
||||
Vendored
+45
-1
@@ -39,12 +39,56 @@ DUMMY_VIDEO_INFO = {
|
||||
"video.crf": 30,
|
||||
"video.preset": 12,
|
||||
"video.fast_decode": 0,
|
||||
"video.is_depth_map": False,
|
||||
"is_depth_map": False,
|
||||
"has_audio": False,
|
||||
}
|
||||
DUMMY_CAMERA_FEATURES = {
|
||||
"laptop": {"shape": (64, 96, 3), "names": ["height", "width", "channels"], "info": DUMMY_VIDEO_INFO},
|
||||
"phone": {"shape": (64, 96, 3), "names": ["height", "width", "channels"], "info": DUMMY_VIDEO_INFO},
|
||||
}
|
||||
DUMMY_DEPTH_VIDEO_INFO = {
|
||||
**DUMMY_VIDEO_INFO,
|
||||
"is_depth_map": True,
|
||||
}
|
||||
DUMMY_DEPTH_VIDEO_INFO_FULL = {
|
||||
**{k: v for k, v in DUMMY_VIDEO_INFO.items() if k != "video.preset"},
|
||||
"video.codec": "hevc",
|
||||
"video.pix_fmt": "gray12le",
|
||||
"is_depth_map": True,
|
||||
"video.depth_min": 0.05,
|
||||
"video.depth_max": 8.0,
|
||||
"video.shift": 2.5,
|
||||
"video.use_log": True,
|
||||
}
|
||||
DUMMY_DEPTH_CAMERA_FEATURES = {
|
||||
"laptop_depth": {
|
||||
"shape": (64, 96, 1),
|
||||
"names": ["height", "width", "channels"],
|
||||
"info": DUMMY_DEPTH_VIDEO_INFO,
|
||||
},
|
||||
}
|
||||
DUMMY_CAMERA_FEATURES_WITH_DEPTH = {**DUMMY_CAMERA_FEATURES, **DUMMY_DEPTH_CAMERA_FEATURES}
|
||||
DUMMY_CHW = (3, 96, 128)
|
||||
DUMMY_HWC = (96, 128, 3)
|
||||
|
||||
# Default video feature set used by video-encoding persistence tests.
|
||||
DUMMY_VIDEO_FEATURES = {
|
||||
"observation.images.cam": {
|
||||
"dtype": "video",
|
||||
"shape": (64, 96, 3),
|
||||
"names": ["height", "width", "channels"],
|
||||
},
|
||||
"action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]},
|
||||
}
|
||||
DUMMY_VIDEO_KEY = "observation.images.cam"
|
||||
|
||||
DUMMY_DEPTH_FEATURES = {
|
||||
"observation.images.depth": {
|
||||
"dtype": "video",
|
||||
"shape": (64, 96, 1),
|
||||
"names": ["height", "width", "channels"],
|
||||
"info": {"is_depth_map": True},
|
||||
},
|
||||
"action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]},
|
||||
}
|
||||
DUMMY_DEPTH_KEY = "observation.images.depth"
|
||||
|
||||
Vendored
+38
@@ -49,6 +49,39 @@ from tests.fixtures.constants import (
|
||||
)
|
||||
|
||||
|
||||
def add_frames(dataset: LeRobotDataset, num_frames: int) -> None:
|
||||
"""Append ``num_frames`` synthetic frames to ``dataset``.
|
||||
|
||||
Generates per-feature payloads from ``dataset.meta``: uint16 depth ramps for
|
||||
keys in ``dataset.meta.depth_keys``, uint8 random noise for video/image keys,
|
||||
and float32 zeros for everything else. ``DEFAULT_FEATURES`` (timestamp,
|
||||
frame_index, ...) are auto-populated by ``add_frame`` and skipped here.
|
||||
"""
|
||||
video_keys = dataset.meta.video_keys
|
||||
depth_keys = dataset.meta.depth_keys
|
||||
# Smooth gradient base reused per (H, W) to keep depth frames cheap to
|
||||
# encode (HEVC Main 12 hates white noise).
|
||||
_depth_base_cache: dict[tuple[int, int], np.ndarray] = {}
|
||||
for i in range(num_frames):
|
||||
frame: dict = {"task": "test"}
|
||||
for key, ft in dataset.meta.features.items():
|
||||
if key in DEFAULT_FEATURES:
|
||||
continue
|
||||
shape = ft["shape"]
|
||||
if key in depth_keys:
|
||||
h, w, _ = shape
|
||||
base = _depth_base_cache.setdefault(
|
||||
(h, w),
|
||||
np.linspace(100.0, 10_000.0, h * w, dtype=np.float32).reshape(h, w, 1),
|
||||
)
|
||||
frame[key] = (base + 50.0 * i).clip(0, 65535).astype(np.uint16)
|
||||
elif key in video_keys:
|
||||
frame[key] = np.random.randint(0, 256, shape, dtype=np.uint8)
|
||||
else:
|
||||
frame[key] = np.zeros(shape, dtype=np.float32)
|
||||
dataset.add_frame(frame)
|
||||
|
||||
|
||||
class LeRobotDatasetFactory(Protocol):
|
||||
def __call__(self, *args, **kwargs) -> LeRobotDataset: ...
|
||||
|
||||
@@ -485,10 +518,14 @@ def lerobot_dataset_factory(
|
||||
hf_dataset: datasets.Dataset | None = None,
|
||||
data_files_size_in_mb: float = DEFAULT_DATA_FILE_SIZE_IN_MB,
|
||||
chunks_size: int = DEFAULT_CHUNK_SIZE,
|
||||
camera_features: dict | None = None,
|
||||
**kwargs,
|
||||
) -> LeRobotDataset:
|
||||
# Instantiate objects
|
||||
if info is None:
|
||||
info_kwargs = {}
|
||||
if camera_features is not None:
|
||||
info_kwargs["camera_features"] = camera_features
|
||||
info = info_factory(
|
||||
total_episodes=total_episodes,
|
||||
total_frames=total_frames,
|
||||
@@ -496,6 +533,7 @@ def lerobot_dataset_factory(
|
||||
use_videos=use_videos,
|
||||
data_files_size_in_mb=data_files_size_in_mb,
|
||||
chunks_size=chunks_size,
|
||||
**info_kwargs,
|
||||
)
|
||||
if stats is None:
|
||||
stats = stats_factory(features=info.features)
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
# 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.
|
||||
|
||||
# Importing concrete policy configs registers their draccus `--policy.type`
|
||||
# choices (e.g. "act") so tests can parse them.
|
||||
from lerobot.policies.act.configuration_act import ACTConfig # noqa: F401
|
||||
@@ -0,0 +1,66 @@
|
||||
# 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.
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
||||
|
||||
from lerobot.jobs.dataset import ensure_dataset_available
|
||||
|
||||
|
||||
def _api_with_dataset(exists: bool):
|
||||
api = MagicMock()
|
||||
api.repo_exists.return_value = exists
|
||||
return api
|
||||
|
||||
|
||||
def _make_local_cache(tmp_path, repo_id: str) -> None:
|
||||
"""Create the minimal local-cache layout that ensure_dataset_available checks."""
|
||||
info = tmp_path / repo_id / "meta" / "info.json"
|
||||
info.parent.mkdir(parents=True)
|
||||
info.write_text("{}")
|
||||
|
||||
|
||||
# Branch 1: dataset already on Hub → no push, no error (pod downloads by repo_id).
|
||||
def test_dataset_already_on_hub_is_noop():
|
||||
api = _api_with_dataset(True)
|
||||
assert ensure_dataset_available("user/ds", api=api) is None
|
||||
api.repo_exists.assert_called_once_with("user/ds", repo_type="dataset")
|
||||
|
||||
|
||||
# Branch 2: not on Hub but present locally → always push privately.
|
||||
def test_dataset_local_only_uploads_privately(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr("lerobot.jobs.dataset.HF_LEROBOT_HOME", tmp_path)
|
||||
_make_local_cache(tmp_path, "user/ds")
|
||||
|
||||
api = _api_with_dataset(False)
|
||||
mock_ds_cls = MagicMock()
|
||||
monkeypatch.setattr("lerobot.jobs.dataset.LeRobotDataset", mock_ds_cls)
|
||||
|
||||
assert ensure_dataset_available("user/ds", api=api, tags=["lerobot", "lelab"]) is None
|
||||
|
||||
mock_ds_cls.assert_called_once_with("user/ds")
|
||||
mock_ds_cls.return_value.push_to_hub.assert_called_once_with(private=True, tags=["lerobot", "lelab"])
|
||||
|
||||
|
||||
# Branch 3: not on Hub, NOT in local cache → RuntimeError.
|
||||
def test_dataset_neither_on_hub_nor_local_raises(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr("lerobot.jobs.dataset.HF_LEROBOT_HOME", tmp_path)
|
||||
# tmp_path is empty — no local cache.
|
||||
|
||||
api = _api_with_dataset(False)
|
||||
with pytest.raises(RuntimeError, match="not in the local cache"):
|
||||
ensure_dataset_available("user/ds", api=api)
|
||||
@@ -0,0 +1,493 @@
|
||||
# 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.
|
||||
|
||||
import datetime as dt
|
||||
import json
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
|
||||
import draccus
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
||||
|
||||
from lerobot.configs.train import TrainPipelineConfig
|
||||
from lerobot.jobs.hf import (
|
||||
_pod_forwarded_args,
|
||||
_poll_until_done,
|
||||
build_remote_config_file,
|
||||
build_repo_id,
|
||||
resolve_job_tags,
|
||||
resolve_wandb_api_key,
|
||||
submit_to_hf,
|
||||
)
|
||||
|
||||
|
||||
def test_resolve_job_tags_always_includes_lerobot_and_dedups():
|
||||
assert resolve_job_tags(None) == ["lerobot"]
|
||||
assert resolve_job_tags([]) == ["lerobot"]
|
||||
assert resolve_job_tags(["lelab"]) == ["lerobot", "lelab"]
|
||||
# lerobot isn't duplicated if passed explicitly; order is stable.
|
||||
assert resolve_job_tags(["lelab", "lerobot", "lelab"]) == ["lerobot", "lelab"]
|
||||
|
||||
|
||||
def _fake_inspect(stage_value, *, as_enum=True):
|
||||
# huggingface_hub returns `stage` as an enum (with `.value`) in some versions and a plain str in others.
|
||||
stage = SimpleNamespace(value=stage_value) if as_enum else stage_value
|
||||
return lambda job_id: SimpleNamespace(status=SimpleNamespace(stage=stage))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("as_enum", [True, False], ids=["enum_stage", "str_stage"])
|
||||
def test_poll_until_done_returns_terminal_stage(monkeypatch, as_enum):
|
||||
monkeypatch.setattr("lerobot.jobs.hf.inspect_job", _fake_inspect("COMPLETED", as_enum=as_enum))
|
||||
done = threading.Event()
|
||||
assert _poll_until_done("j", done, poll_interval=0.01) == "COMPLETED"
|
||||
assert done.is_set()
|
||||
|
||||
|
||||
def test_poll_until_done_exits_when_done_already_set(monkeypatch):
|
||||
# Non-terminal forever; with done pre-set the loop must not block and returns None.
|
||||
monkeypatch.setattr("lerobot.jobs.hf.inspect_job", _fake_inspect("RUNNING"))
|
||||
done = threading.Event()
|
||||
done.set()
|
||||
assert _poll_until_done("j", done, poll_interval=0.01) is None
|
||||
|
||||
|
||||
def test_poll_until_done_gives_up_after_repeated_network_failures(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"lerobot.jobs.hf.inspect_job", lambda job_id: (_ for _ in ()).throw(httpx.ConnectError("boom"))
|
||||
)
|
||||
done = threading.Event()
|
||||
result = _poll_until_done("j", done, poll_interval=0.001, max_failures=3)
|
||||
assert result is None
|
||||
assert done.is_set()
|
||||
|
||||
|
||||
def test_poll_until_done_propagates_programming_errors(monkeypatch):
|
||||
"""A bug (e.g. TypeError) must surface, not be silently retried as a transient failure."""
|
||||
monkeypatch.setattr("lerobot.jobs.hf.inspect_job", lambda job_id: (_ for _ in ()).throw(TypeError("bug")))
|
||||
done = threading.Event()
|
||||
with pytest.raises(TypeError):
|
||||
_poll_until_done("j", done, poll_interval=0.001, max_failures=3)
|
||||
|
||||
|
||||
def test_resolve_wandb_key_from_env(monkeypatch):
|
||||
monkeypatch.setenv("WANDB_API_KEY", "abc123")
|
||||
assert resolve_wandb_api_key() == "abc123"
|
||||
|
||||
|
||||
def test_resolve_wandb_key_missing(monkeypatch, tmp_path):
|
||||
monkeypatch.delenv("WANDB_API_KEY", raising=False)
|
||||
monkeypatch.setenv("HOME", str(tmp_path)) # no ~/.netrc here
|
||||
monkeypatch.setattr("netrc.netrc", lambda *a, **k: (_ for _ in ()).throw(FileNotFoundError()))
|
||||
assert resolve_wandb_api_key() is None
|
||||
|
||||
|
||||
def test_resolve_wandb_key_from_netrc(monkeypatch):
|
||||
# No env var → fall back to the wandb credentials in ~/.netrc.
|
||||
monkeypatch.delenv("WANDB_API_KEY", raising=False)
|
||||
|
||||
class _FakeNetrc:
|
||||
def authenticators(self, host):
|
||||
assert host == "api.wandb.ai"
|
||||
return ("login", "account", "netrc-secret")
|
||||
|
||||
monkeypatch.setattr("netrc.netrc", lambda *a, **k: _FakeNetrc())
|
||||
assert resolve_wandb_api_key() == "netrc-secret"
|
||||
|
||||
|
||||
def test_resolve_wandb_key_netrc_without_wandb_entry(monkeypatch):
|
||||
# ~/.netrc exists but has no api.wandb.ai entry → None.
|
||||
monkeypatch.delenv("WANDB_API_KEY", raising=False)
|
||||
|
||||
class _FakeNetrc:
|
||||
def authenticators(self, host):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr("netrc.netrc", lambda *a, **k: _FakeNetrc())
|
||||
assert resolve_wandb_api_key() is None
|
||||
|
||||
|
||||
def test_build_repo_id_sanitizes_and_timestamps():
|
||||
now = dt.datetime(2026, 6, 19, 10, 22, 3)
|
||||
assert build_repo_id("alice", "act", now) == "alice/act_2026-06-19_10-22-03"
|
||||
# Runs of illegal characters collapse to a single dash; edges are trimmed.
|
||||
assert build_repo_id("alice", "my cool/run!!", now) == "alice/my-cool-run_2026-06-19_10-22-03"
|
||||
# A name with nothing usable falls back to "train".
|
||||
assert build_repo_id("alice", "///", now) == "alice/train_2026-06-19_10-22-03"
|
||||
|
||||
|
||||
def test_pod_forwarded_args_drops_host_only_flags():
|
||||
"""User overrides are replayed on the pod, minus flags that only make sense on the submitter.
|
||||
|
||||
`--dataset.root` is a host-local path the pod can't read, so it must be dropped in both the
|
||||
`--name=value` and `--name value` forms; unrelated overrides are forwarded untouched.
|
||||
"""
|
||||
argv = [
|
||||
"--config_path=u/d",
|
||||
"--dataset.root=/local/data",
|
||||
"--dataset.root",
|
||||
"/other/local/data",
|
||||
"--policy.repo_id=u/keep",
|
||||
"--steps=10",
|
||||
"--job.target=a10g-small",
|
||||
]
|
||||
forwarded = _pod_forwarded_args(
|
||||
argv,
|
||||
drop_names=("--config_path", "--policy.repo_id", "--policy.push_to_hub", "--dataset.root"),
|
||||
drop_prefixes=("--job.",),
|
||||
)
|
||||
assert forwarded == ["--steps=10"]
|
||||
|
||||
|
||||
def _minimal_cfg():
|
||||
return draccus.parse(
|
||||
TrainPipelineConfig,
|
||||
args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"],
|
||||
)
|
||||
|
||||
|
||||
def test_validate_skips_repo_id_check_for_remote():
|
||||
"""Remote runs auto-assign repo_id in submit_to_hf, so validate() must not demand it up front."""
|
||||
cfg = _minimal_cfg() # remote target, push_to_hub default True, no explicit repo_id
|
||||
assert cfg.policy.repo_id is None
|
||||
cfg.validate() # must not raise
|
||||
|
||||
|
||||
def test_validate_requires_repo_id_for_local_push():
|
||||
"""Local runs that push to the Hub still need an explicit repo_id."""
|
||||
cfg = draccus.parse(
|
||||
TrainPipelineConfig,
|
||||
args=["--dataset.repo_id", "u/d", "--policy.type", "act"],
|
||||
)
|
||||
with pytest.raises(ValueError, match="repo_id"):
|
||||
cfg.validate()
|
||||
|
||||
|
||||
def test_build_remote_config_applies_overrides(tmp_path):
|
||||
cfg = _minimal_cfg()
|
||||
dest = tmp_path / "train_config.json"
|
||||
out = build_remote_config_file(cfg, "u/run", dest)
|
||||
assert out == dest
|
||||
data = json.loads(dest.read_text())
|
||||
# `job` is client-only orchestration and must be stripped for the pod.
|
||||
assert "job" not in data
|
||||
# save_checkpoint_to_hub defaults off → omitted so older images accept the config.
|
||||
assert "save_checkpoint_to_hub" not in data
|
||||
assert data["policy"]["push_to_hub"] is True
|
||||
assert data["policy"]["repo_id"] == "u/run"
|
||||
assert data["policy"]["device"] is None # pod auto-detects its GPU
|
||||
assert data["dataset"]["root"] is None # pod resolves the dataset by repo_id
|
||||
# the caller's cfg must be left untouched (function works on a deep copy)
|
||||
assert cfg.job.target == "a10g-small"
|
||||
assert cfg.save_checkpoint_to_hub is False
|
||||
|
||||
|
||||
def test_build_remote_config_includes_checkpoint_flag_when_enabled(tmp_path):
|
||||
cfg = draccus.parse(
|
||||
TrainPipelineConfig,
|
||||
args=[
|
||||
"--dataset.repo_id",
|
||||
"u/d",
|
||||
"--policy.type",
|
||||
"act",
|
||||
"--job.target",
|
||||
"a10g-small",
|
||||
"--save_checkpoint_to_hub",
|
||||
"true",
|
||||
],
|
||||
)
|
||||
dest = tmp_path / "train_config.json"
|
||||
build_remote_config_file(cfg, "u/run", dest)
|
||||
data = json.loads(dest.read_text())
|
||||
# explicitly enabled → kept in the config (requires a matching trainer image).
|
||||
assert data["save_checkpoint_to_hub"] is True
|
||||
assert "job" not in data
|
||||
|
||||
|
||||
def test_build_remote_config_merges_tags_into_policy(tmp_path):
|
||||
cfg = _minimal_cfg()
|
||||
dest = tmp_path / "train_config.json"
|
||||
build_remote_config_file(cfg, "u/run", dest, tags=["lerobot", "lelab"])
|
||||
data = json.loads(dest.read_text())
|
||||
# tags propagate to the model the pod pushes.
|
||||
assert data["policy"]["tags"] == ["lerobot", "lelab"]
|
||||
|
||||
|
||||
def test_build_remote_config_merges_tags_without_duplicating(tmp_path):
|
||||
cfg = _minimal_cfg()
|
||||
cfg.policy.tags = ["existing", "lerobot"]
|
||||
dest = tmp_path / "train_config.json"
|
||||
build_remote_config_file(cfg, "u/run", dest, tags=["lerobot", "lelab"])
|
||||
data = json.loads(dest.read_text())
|
||||
# pre-existing policy tags are kept; only genuinely-new tags are appended (no dup "lerobot").
|
||||
assert data["policy"]["tags"] == ["existing", "lerobot", "lelab"]
|
||||
|
||||
|
||||
def test_submit_requires_login(monkeypatch):
|
||||
monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: None)
|
||||
cfg = draccus.parse(
|
||||
TrainPipelineConfig,
|
||||
args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"],
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="hf auth login"):
|
||||
submit_to_hf(cfg)
|
||||
|
||||
|
||||
def test_submit_passes_validation_and_submits(monkeypatch):
|
||||
"""A type-based policy with no explicit repo_id is auto-assigned one and submitted."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
# Patch get_token
|
||||
monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok")
|
||||
|
||||
# Patch HfApi so whoami returns alice
|
||||
class FakeHfApi:
|
||||
def __init__(self, token=None):
|
||||
pass
|
||||
|
||||
def whoami(self, token=None):
|
||||
return {"name": "alice"}
|
||||
|
||||
monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi)
|
||||
|
||||
# ensure_dataset_available returns None; patch it out so no Hub access happens
|
||||
# (hf.py imports it at module level, so patch it on lerobot.jobs.hf).
|
||||
monkeypatch.setattr("lerobot.jobs.hf.ensure_dataset_available", lambda *a, **kw: None)
|
||||
|
||||
# Patch _stage_config_on_hub to skip network
|
||||
monkeypatch.setattr(
|
||||
"lerobot.jobs.hf._stage_config_on_hub",
|
||||
lambda cfg, repo_id, token, tags=None: repo_id,
|
||||
)
|
||||
|
||||
# Patch run_job to return a fake job
|
||||
fake_job = MagicMock()
|
||||
fake_job.id = "job-123"
|
||||
run_job_calls = []
|
||||
|
||||
def fake_run_job(**kwargs):
|
||||
run_job_calls.append(kwargs)
|
||||
return fake_job
|
||||
|
||||
monkeypatch.setattr("lerobot.jobs.hf.run_job", fake_run_job)
|
||||
|
||||
cfg = draccus.parse(
|
||||
TrainPipelineConfig,
|
||||
args=[
|
||||
"--dataset.repo_id",
|
||||
"u/d",
|
||||
"--policy.type",
|
||||
"act",
|
||||
"--job.target",
|
||||
"a10g-small",
|
||||
"--job.detach",
|
||||
"true",
|
||||
],
|
||||
)
|
||||
|
||||
# Must NOT raise (pre-fix this raised ValueError about missing repo_id)
|
||||
submit_to_hf(cfg)
|
||||
|
||||
assert len(run_job_calls) == 1, "run_job should have been called exactly once"
|
||||
assert cfg.policy.repo_id is not None
|
||||
assert cfg.policy.repo_id.startswith("alice/")
|
||||
call = run_job_calls[0]
|
||||
# The pod runs `lerobot-train --config_path=<staged repo>` on the requested flavor/image.
|
||||
assert call["command"][0] == "lerobot-train"
|
||||
assert call["command"][1].startswith("--config_path=")
|
||||
assert call["flavor"] == "a10g-small"
|
||||
assert call["image"] == "huggingface/lerobot-gpu:latest"
|
||||
# The Hub token is forwarded so the pod can pull the (possibly private) dataset.
|
||||
assert call["secrets"]["HF_TOKEN"] == "tok"
|
||||
# Every job carries the lerobot tag as a queryable label.
|
||||
assert call["labels"].get("lerobot") == "true"
|
||||
|
||||
|
||||
def test_submit_rejects_reward_model_training(monkeypatch):
|
||||
"""Remote training only supports policies; reward-model runs fail fast with a clear error."""
|
||||
monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok")
|
||||
|
||||
class FakeHfApi:
|
||||
def __init__(self, token=None):
|
||||
pass
|
||||
|
||||
def whoami(self, token=None):
|
||||
return {"name": "alice"}
|
||||
|
||||
monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi)
|
||||
|
||||
cfg = _minimal_cfg()
|
||||
cfg.reward_model = SimpleNamespace(type="reward") # marks this as reward-model training
|
||||
monkeypatch.setattr(cfg, "validate", lambda: None) # skip pretrained-path resolution
|
||||
|
||||
with pytest.raises(ValueError, match="reward model"):
|
||||
submit_to_hf(cfg)
|
||||
|
||||
|
||||
@pytest.mark.timeout(15)
|
||||
def test_submit_returns_when_job_completes(monkeypatch):
|
||||
"""Non-detach path must RETURN (not hang) once the job reaches a terminal stage."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok")
|
||||
|
||||
class FakeHfApi:
|
||||
def __init__(self, token=None):
|
||||
pass
|
||||
|
||||
def whoami(self, token=None):
|
||||
return {"name": "alice"}
|
||||
|
||||
monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi)
|
||||
monkeypatch.setattr("lerobot.jobs.hf.ensure_dataset_available", lambda *a, **kw: None)
|
||||
monkeypatch.setattr(
|
||||
"lerobot.jobs.hf._stage_config_on_hub", lambda cfg, repo_id, token, tags=None: repo_id
|
||||
)
|
||||
monkeypatch.setattr("lerobot.jobs.hf.run_job", lambda **kw: SimpleNamespace(id="job-1", url="http://x"))
|
||||
# Job is already COMPLETED on the first poll.
|
||||
monkeypatch.setattr(
|
||||
"lerobot.jobs.hf.inspect_job",
|
||||
lambda job_id: SimpleNamespace(
|
||||
status=SimpleNamespace(stage=SimpleNamespace(value="COMPLETED"), message=None)
|
||||
),
|
||||
)
|
||||
# Log stream ends immediately.
|
||||
monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter(()))
|
||||
|
||||
cfg = draccus.parse(
|
||||
TrainPipelineConfig,
|
||||
args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"],
|
||||
)
|
||||
# Runs in the pytest main thread (signal handler install requires it); the
|
||||
# @timeout marker fails the test instead of hanging if it regresses.
|
||||
submit_to_hf(cfg)
|
||||
|
||||
|
||||
@pytest.mark.timeout(15)
|
||||
def test_submit_returns_on_model_pushed_marker(monkeypatch):
|
||||
"""Finish when the model-pushed log appears, even if the job stage never flips."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok")
|
||||
|
||||
class FakeHfApi:
|
||||
def __init__(self, token=None):
|
||||
pass
|
||||
|
||||
def whoami(self, token=None):
|
||||
return {"name": "alice"}
|
||||
|
||||
monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi)
|
||||
monkeypatch.setattr("lerobot.jobs.hf.ensure_dataset_available", lambda *a, **kw: None)
|
||||
monkeypatch.setattr(
|
||||
"lerobot.jobs.hf._stage_config_on_hub", lambda cfg, repo_id, token, tags=None: repo_id
|
||||
)
|
||||
monkeypatch.setattr("lerobot.jobs.hf.run_job", lambda **kw: SimpleNamespace(id="job-1", url="http://x"))
|
||||
# Job stays RUNNING forever — only the log marker can end the command.
|
||||
monkeypatch.setattr(
|
||||
"lerobot.jobs.hf.inspect_job",
|
||||
lambda job_id: SimpleNamespace(
|
||||
status=SimpleNamespace(stage=SimpleNamespace(value="RUNNING"), message=None)
|
||||
),
|
||||
)
|
||||
pushed_line = "INFO Model pushed to https://huggingface.co/alice/myrun"
|
||||
monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter([pushed_line]))
|
||||
|
||||
cfg = draccus.parse(
|
||||
TrainPipelineConfig,
|
||||
args=[
|
||||
"--dataset.repo_id",
|
||||
"u/d",
|
||||
"--policy.type",
|
||||
"act",
|
||||
"--policy.repo_id",
|
||||
"alice/myrun",
|
||||
"--job.target",
|
||||
"a10g-small",
|
||||
],
|
||||
)
|
||||
# Must return via the model-pushed marker despite the perpetual RUNNING stage.
|
||||
submit_to_hf(cfg)
|
||||
|
||||
|
||||
def test_submit_raises_when_wandb_enabled_without_key(monkeypatch):
|
||||
"""wandb.enable with no key reachable anywhere fails fast, before submitting."""
|
||||
|
||||
monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok")
|
||||
|
||||
class FakeHfApi:
|
||||
def __init__(self, token=None):
|
||||
pass
|
||||
|
||||
def whoami(self, token=None):
|
||||
return {"name": "alice"}
|
||||
|
||||
monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi)
|
||||
monkeypatch.setattr("lerobot.jobs.hf.resolve_wandb_api_key", lambda: None)
|
||||
|
||||
cfg = draccus.parse(
|
||||
TrainPipelineConfig,
|
||||
args=[
|
||||
"--dataset.repo_id",
|
||||
"u/d",
|
||||
"--policy.type",
|
||||
"act",
|
||||
"--job.target",
|
||||
"a10g-small",
|
||||
"--wandb.enable",
|
||||
"true",
|
||||
],
|
||||
)
|
||||
with pytest.raises(ValueError, match="WANDB_API_KEY"):
|
||||
submit_to_hf(cfg)
|
||||
|
||||
|
||||
@pytest.mark.timeout(15)
|
||||
def test_submit_raises_when_job_ends_in_error(monkeypatch):
|
||||
"""A terminal non-COMPLETED stage with no model-pushed marker must raise with the status."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok")
|
||||
|
||||
class FakeHfApi:
|
||||
def __init__(self, token=None):
|
||||
pass
|
||||
|
||||
def whoami(self, token=None):
|
||||
return {"name": "alice"}
|
||||
|
||||
monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi)
|
||||
monkeypatch.setattr("lerobot.jobs.hf.ensure_dataset_available", lambda *a, **kw: None)
|
||||
monkeypatch.setattr(
|
||||
"lerobot.jobs.hf._stage_config_on_hub", lambda cfg, repo_id, token, tags=None: repo_id
|
||||
)
|
||||
monkeypatch.setattr("lerobot.jobs.hf.run_job", lambda **kw: SimpleNamespace(id="job-1", url="http://x"))
|
||||
# Job fails: a terminal ERROR stage carrying the platform's status message.
|
||||
monkeypatch.setattr(
|
||||
"lerobot.jobs.hf.inspect_job",
|
||||
lambda job_id: SimpleNamespace(
|
||||
status=SimpleNamespace(stage=SimpleNamespace(value="ERROR"), message="Job timeout")
|
||||
),
|
||||
)
|
||||
# Logs end without the model-pushed marker.
|
||||
monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter(()))
|
||||
|
||||
cfg = draccus.parse(
|
||||
TrainPipelineConfig,
|
||||
args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"],
|
||||
)
|
||||
with pytest.raises(RuntimeError, match=r"stage=ERROR \(Job timeout\)"):
|
||||
submit_to_hf(cfg)
|
||||
@@ -0,0 +1,64 @@
|
||||
# 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.
|
||||
|
||||
import draccus
|
||||
import pytest
|
||||
|
||||
from lerobot.configs import JobConfig
|
||||
from lerobot.configs.train import TrainPipelineConfig
|
||||
|
||||
|
||||
def test_jobconfig_defaults_are_local():
|
||||
cfg = JobConfig()
|
||||
assert cfg.target is None
|
||||
assert cfg.is_remote is False
|
||||
assert cfg.image == "huggingface/lerobot-gpu:latest"
|
||||
assert cfg.timeout == "2d"
|
||||
assert cfg.detach is False
|
||||
|
||||
|
||||
def test_jobconfig_local_string_is_not_remote():
|
||||
assert JobConfig(target="local").is_remote is False
|
||||
|
||||
|
||||
def test_jobconfig_flavor_is_remote():
|
||||
assert JobConfig(target="a10g-small").is_remote is True
|
||||
|
||||
|
||||
def test_train_config_parses_job_target():
|
||||
parsed = draccus.parse(
|
||||
TrainPipelineConfig,
|
||||
args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"],
|
||||
)
|
||||
assert parsed.job.target == "a10g-small"
|
||||
assert parsed.job.is_remote is True
|
||||
assert parsed.save_checkpoint_to_hub is False
|
||||
|
||||
|
||||
def test_save_checkpoint_to_hub_requires_repo_id():
|
||||
cfg = draccus.parse(
|
||||
TrainPipelineConfig,
|
||||
args=[
|
||||
"--dataset.repo_id",
|
||||
"u/d",
|
||||
"--policy.type",
|
||||
"act",
|
||||
"--policy.push_to_hub",
|
||||
"false",
|
||||
"--save_checkpoint_to_hub",
|
||||
"true",
|
||||
],
|
||||
)
|
||||
with pytest.raises(ValueError, match="requires --policy.repo_id"):
|
||||
cfg.validate()
|
||||
@@ -20,6 +20,7 @@ from lerobot.optim.optimizers import (
|
||||
MultiAdamConfig,
|
||||
SGDConfig,
|
||||
load_optimizer_state,
|
||||
load_optimizer_state_dict,
|
||||
save_optimizer_state,
|
||||
)
|
||||
from lerobot.utils.constants import (
|
||||
@@ -65,6 +66,44 @@ def test_save_and_load_optimizer_state(model_params, optimizer, tmp_path):
|
||||
torch.testing.assert_close(optimizer.state_dict(), loaded_optimizer.state_dict())
|
||||
|
||||
|
||||
def test_save_and_load_fsdp_optimizer_state_dict_roundtrip(tmp_path):
|
||||
"""The FSDP full optimizer state dict is keyed by parameter FQNs (dotted strings), not the
|
||||
integer indices of the single-GPU path. Verify it survives the safetensors save -> read
|
||||
round-trip used by the FSDP save/resume path (save_optimizer_state(optim_state_dict=...) then
|
||||
load_optimizer_state_dict), which the flatten/unflatten "/" separator must not corrupt."""
|
||||
full_osd = {
|
||||
"state": {
|
||||
"model.layers.0.weight": {
|
||||
"step": torch.tensor(3.0),
|
||||
"exp_avg": torch.randn(4, 4),
|
||||
"exp_avg_sq": torch.randn(4, 4),
|
||||
},
|
||||
"model.layers.0.bias": {
|
||||
"step": torch.tensor(3.0),
|
||||
"exp_avg": torch.randn(4),
|
||||
"exp_avg_sq": torch.randn(4),
|
||||
},
|
||||
},
|
||||
"param_groups": [
|
||||
{"lr": 1e-4, "betas": [0.9, 0.999], "eps": 1e-8, "weight_decay": 0.0, "params": [0, 1]}
|
||||
],
|
||||
}
|
||||
|
||||
save_optimizer_state(
|
||||
torch.optim.Adam([torch.nn.Parameter(torch.randn(1))]), tmp_path, optim_state_dict=full_osd
|
||||
)
|
||||
assert (tmp_path / OPTIMIZER_STATE).is_file()
|
||||
assert (tmp_path / OPTIMIZER_PARAM_GROUPS).is_file()
|
||||
|
||||
loaded = load_optimizer_state_dict(tmp_path)
|
||||
# FQN keys must be preserved verbatim (not int-cast, not split on their dots).
|
||||
assert set(loaded["state"].keys()) == set(full_osd["state"].keys())
|
||||
for fqn, sub in full_osd["state"].items():
|
||||
for k, v in sub.items():
|
||||
torch.testing.assert_close(loaded["state"][fqn][k], v)
|
||||
assert loaded["param_groups"] == full_osd["param_groups"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def base_params_dict():
|
||||
return {
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
@@ -35,24 +33,27 @@ pytest.importorskip("scipy")
|
||||
from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature
|
||||
from lerobot.policies import get_policy_class, make_policy_config
|
||||
from lerobot.policies.molmoact2 import (
|
||||
configuration_molmoact2 as molmoact2_config,
|
||||
modeling_molmoact2 as molmoact2_modeling,
|
||||
processor_molmoact2 as molmoact2_processor,
|
||||
)
|
||||
from lerobot.policies.molmoact2.configuration_molmoact2 import (
|
||||
MolmoAct2Config,
|
||||
MolmoAct2CosineDecayWithWarmupSchedulerConfig,
|
||||
infer_molmoact2_max_sequence_length,
|
||||
from lerobot.policies.molmoact2.configuration_molmoact2 import MolmoAct2Config
|
||||
from lerobot.policies.molmoact2.modeling_molmoact2 import (
|
||||
MolmoAct2Policy,
|
||||
_apply_action_chunk_padding_mask,
|
||||
_apply_action_dim_padding_mask,
|
||||
_combine_rollout_seeds,
|
||||
)
|
||||
from lerobot.policies.molmoact2.modeling_molmoact2 import MolmoAct2Policy
|
||||
from lerobot.policies.molmoact2.processor_molmoact2 import (
|
||||
MolmoAct2ActionFrameTransformStep,
|
||||
MolmoAct2ClampNormalizedProcessorStep,
|
||||
MolmoAct2MaskedNormalizerProcessorStep,
|
||||
MolmoAct2MaskedUnnormalizerProcessorStep,
|
||||
MolmoAct2PackInputsProcessorStep,
|
||||
MolmoAct2StateFrameTransformStep,
|
||||
_add_gripper_masks_to_stats,
|
||||
_build_discrete_state_string,
|
||||
_normalize_question_text,
|
||||
infer_molmoact2_max_sequence_length,
|
||||
make_molmoact2_pre_post_processors,
|
||||
)
|
||||
from lerobot.policies.rtc.configuration_rtc import RTCConfig
|
||||
@@ -71,34 +72,38 @@ def test_molmoact2_policy_registration():
|
||||
assert cfg.per_episode_seed is False
|
||||
assert cfg.eval_seed is None
|
||||
assert cfg.normalize_language is True
|
||||
assert cfg.get_scheduler_preset().num_decay_steps is None
|
||||
assert cfg.get_scheduler_preset().num_decay_steps == 100_000
|
||||
assert cfg.action_delta_indices == list(range(cfg.chunk_size))
|
||||
assert get_policy_class("molmoact2") is MolmoAct2Policy
|
||||
|
||||
|
||||
def test_molmoact2_checkpoint_download_ignores_remote_python(monkeypatch):
|
||||
import huggingface_hub
|
||||
|
||||
download_kwargs = {}
|
||||
|
||||
def fake_snapshot_download(**kwargs):
|
||||
download_kwargs.update(kwargs)
|
||||
return "/tmp/downloaded-molmoact2"
|
||||
|
||||
monkeypatch.setattr(molmoact2_config, "snapshot_download", fake_snapshot_download)
|
||||
monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
|
||||
|
||||
checkpoint_location = molmoact2_config._resolve_checkpoint_location("allenai/MolmoAct2")
|
||||
checkpoint_location = molmoact2_modeling._resolve_checkpoint_location("allenai/MolmoAct2")
|
||||
|
||||
assert checkpoint_location == "/tmp/downloaded-molmoact2"
|
||||
assert download_kwargs["ignore_patterns"] == ["*.py", "*.pyc", "__pycache__/*"]
|
||||
|
||||
|
||||
def test_molmoact2_scheduler_decay_steps_auto_match_training_steps():
|
||||
def test_molmoact2_scheduler_auto_scales_to_training_steps():
|
||||
from lerobot.optim import CosineDecayWithWarmupSchedulerConfig
|
||||
|
||||
param = torch.nn.Parameter(torch.ones(()))
|
||||
optimizer = torch.optim.AdamW([param], lr=0.001)
|
||||
config = MolmoAct2CosineDecayWithWarmupSchedulerConfig(
|
||||
config = CosineDecayWithWarmupSchedulerConfig(
|
||||
peak_lr=0.01,
|
||||
decay_lr=0.001,
|
||||
num_warmup_steps=10,
|
||||
num_decay_steps=None,
|
||||
num_decay_steps=100_000,
|
||||
)
|
||||
|
||||
scheduler = config.build(optimizer, num_training_steps=100)
|
||||
@@ -123,9 +128,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task():
|
||||
batch_size=3,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
expected_first = torch.Generator().manual_seed(
|
||||
MolmoAct2Policy._combine_rollout_seeds(first_seed=1000, batch_size=3)
|
||||
)
|
||||
expected_first = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1000, batch_size=3))
|
||||
assert torch.allclose(torch.rand(4, generator=first), torch.rand(4, generator=expected_first))
|
||||
|
||||
policy.reset()
|
||||
@@ -134,9 +137,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task():
|
||||
batch_size=3,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
expected_second = torch.Generator().manual_seed(
|
||||
MolmoAct2Policy._combine_rollout_seeds(first_seed=1003, batch_size=3)
|
||||
)
|
||||
expected_second = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1003, batch_size=3))
|
||||
assert torch.allclose(torch.rand(4, generator=second), torch.rand(4, generator=expected_second))
|
||||
|
||||
policy.reset()
|
||||
@@ -145,9 +146,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task():
|
||||
batch_size=3,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
expected_new_task = torch.Generator().manual_seed(
|
||||
MolmoAct2Policy._combine_rollout_seeds(first_seed=1000, batch_size=3)
|
||||
)
|
||||
expected_new_task = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1000, batch_size=3))
|
||||
assert torch.allclose(torch.rand(4, generator=new_task), torch.rand(4, generator=expected_new_task))
|
||||
|
||||
|
||||
@@ -537,36 +536,26 @@ def test_train_action_expert_only_requires_continuous_action_mode():
|
||||
|
||||
|
||||
def test_molmoact2_sequence_length_is_inferred_from_fixed_token_budget():
|
||||
cfg = MolmoAct2Config(
|
||||
action_mode="both",
|
||||
chunk_size=10,
|
||||
n_action_steps=10,
|
||||
image_keys=["observation.images.image", "observation.images.wrist_image"],
|
||||
input_features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,))},
|
||||
output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,))},
|
||||
)
|
||||
|
||||
assert cfg.max_sequence_length is None
|
||||
assert cfg.inferred_max_sequence_length() == 640
|
||||
assert cfg.inferred_max_sequence_length(include_discrete_action=False) == 576
|
||||
assert (
|
||||
infer_molmoact2_max_sequence_length(
|
||||
num_images=2,
|
||||
state_dim=8,
|
||||
action_dim=7,
|
||||
action_horizon=30,
|
||||
include_discrete_action=True,
|
||||
num_images=2, state_dim=8, action_dim=7, action_horizon=10, include_discrete_action=True
|
||||
)
|
||||
== 640
|
||||
)
|
||||
assert (
|
||||
infer_molmoact2_max_sequence_length(
|
||||
num_images=2, state_dim=8, action_dim=7, action_horizon=10, include_discrete_action=False
|
||||
)
|
||||
== 576
|
||||
)
|
||||
assert (
|
||||
infer_molmoact2_max_sequence_length(
|
||||
num_images=2, state_dim=8, action_dim=7, action_horizon=30, include_discrete_action=True
|
||||
)
|
||||
== 768
|
||||
)
|
||||
|
||||
|
||||
def test_molmoact2_sequence_length_override_is_preserved():
|
||||
cfg = MolmoAct2Config(max_sequence_length=1024)
|
||||
|
||||
assert cfg.inferred_max_sequence_length(num_images=2, state_dim=8, action_dim=7) == 1024
|
||||
|
||||
|
||||
def test_train_action_expert_only_freezes_non_action_expert_params():
|
||||
class DummyBackbone(torch.nn.Module):
|
||||
def __init__(self):
|
||||
@@ -939,6 +928,39 @@ def test_question_normalization_matches_release_prompt_style():
|
||||
)
|
||||
|
||||
|
||||
def test_joint_frame_transform_round_trip():
|
||||
signs = [1.0, -1.0, 1.0, 1.0, 1.0, 1.0]
|
||||
offsets = [0.0, 90.0, 90.0, 0.0, 0.0, 0.0]
|
||||
original_state = torch.tensor([[10.0, -90.0, -120.0, 30.0, 0.0, -45.0]])
|
||||
|
||||
state_step = MolmoAct2StateFrameTransformStep(joint_signs=signs, joint_offsets=offsets)
|
||||
action_step = MolmoAct2ActionFrameTransformStep(joint_signs=signs, joint_offsets=offsets)
|
||||
|
||||
transition = {
|
||||
TransitionKey.OBSERVATION: {OBS_STATE: original_state.clone()},
|
||||
}
|
||||
transformed = state_step(transition)
|
||||
model_state = transformed[TransitionKey.OBSERVATION][OBS_STATE]
|
||||
|
||||
action_transition = {TransitionKey.ACTION: model_state.clone()}
|
||||
recovered = action_step(action_transition)
|
||||
recovered_state = recovered[TransitionKey.ACTION]
|
||||
|
||||
assert torch.allclose(recovered_state, original_state)
|
||||
|
||||
|
||||
def test_joint_frame_transform_noop_when_none():
|
||||
state_step = MolmoAct2StateFrameTransformStep(joint_signs=None, joint_offsets=None)
|
||||
action_step = MolmoAct2ActionFrameTransformStep(joint_signs=None, joint_offsets=None)
|
||||
state = torch.tensor([[10.0, -90.0, -120.0]])
|
||||
|
||||
state_transition = {TransitionKey.OBSERVATION: {OBS_STATE: state}}
|
||||
assert state_step(state_transition) is state_transition
|
||||
|
||||
action_transition = {TransitionKey.ACTION: state}
|
||||
assert action_step(action_transition) is action_transition
|
||||
|
||||
|
||||
def test_action_padding_marks_only_real_dimensions():
|
||||
step = object.__new__(MolmoAct2PackInputsProcessorStep)
|
||||
step.max_action_dim = 32
|
||||
@@ -963,7 +985,7 @@ def test_action_dim_padding_loss_reduces_like_old_trainer():
|
||||
]
|
||||
)
|
||||
|
||||
reduced = MolmoAct2Policy._apply_action_dim_padding_mask(loss, action_dim_is_pad)
|
||||
reduced = _apply_action_dim_padding_mask(loss, action_dim_is_pad)
|
||||
|
||||
expected = torch.stack(
|
||||
[
|
||||
@@ -979,7 +1001,7 @@ def test_action_chunk_padding_keeps_old_mean_denominator():
|
||||
loss = torch.ones(1, 2, 4, 3)
|
||||
action_horizon_is_pad = torch.tensor([[False, False, True, True]])
|
||||
|
||||
masked = MolmoAct2Policy._apply_action_chunk_padding_mask(loss, action_horizon_is_pad)
|
||||
masked = _apply_action_chunk_padding_mask(loss, action_horizon_is_pad)
|
||||
|
||||
assert masked.mean().item() == 0.5
|
||||
|
||||
|
||||
@@ -23,6 +23,7 @@ import torch
|
||||
|
||||
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
||||
|
||||
from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
|
||||
from packaging import version
|
||||
from safetensors.torch import load_file
|
||||
|
||||
@@ -300,6 +301,29 @@ def test_save_and_load_pretrained(dummy_dataset_metadata, tmp_path, policy_name:
|
||||
torch.testing.assert_close(list(policy.parameters()), list(loaded_policy.parameters()), rtol=0, atol=0)
|
||||
|
||||
|
||||
def test_save_pretrained_with_state_dict(dummy_dataset_metadata, tmp_path):
|
||||
"""Exercise the FSDP checkpoint path: save_pretrained with a pre-gathered state_dict."""
|
||||
policy_cls = get_policy_class("act")
|
||||
policy_cfg = make_policy_config("act")
|
||||
features = dataset_to_policy_features(dummy_dataset_metadata.features)
|
||||
policy_cfg.output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
|
||||
policy_cfg.input_features = {
|
||||
key: ft for key, ft in features.items() if key not in policy_cfg.output_features
|
||||
}
|
||||
policy = policy_cls(policy_cfg)
|
||||
policy.to(policy_cfg.device)
|
||||
|
||||
save_dir = tmp_path / "fsdp_state_dict"
|
||||
policy.save_pretrained(save_dir, state_dict=policy.state_dict())
|
||||
|
||||
# A single, unsharded safetensors file (no sharded set + index).
|
||||
assert (save_dir / SAFETENSORS_SINGLE_FILE).is_file()
|
||||
assert not (save_dir / f"{SAFETENSORS_SINGLE_FILE}.index.json").exists()
|
||||
|
||||
loaded_policy = policy_cls.from_pretrained(save_dir, config=policy_cfg)
|
||||
torch.testing.assert_close(list(policy.parameters()), list(loaded_policy.parameters()), rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("multikey", [True, False])
|
||||
def test_multikey_construction(multikey: bool):
|
||||
"""
|
||||
|
||||
@@ -8,7 +8,6 @@ from types import SimpleNamespace
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torch import Tensor, nn
|
||||
|
||||
from lerobot.configs.types import FeatureType, PolicyFeature
|
||||
@@ -191,7 +190,7 @@ class _FakeQwenInterface(nn.Module):
|
||||
|
||||
def build_inputs(
|
||||
self,
|
||||
images: list[list[Image.Image]],
|
||||
images: list[list[Tensor]],
|
||||
instructions: list[str],
|
||||
action_prompt: str,
|
||||
embodied_prompt: str,
|
||||
@@ -214,12 +213,13 @@ class _FakeQwenInterface(nn.Module):
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def tensor_to_pil(image_tensor: Tensor) -> Image.Image:
|
||||
image = image_tensor.detach().cpu()
|
||||
if image.ndim == 3 and image.shape[0] in (1, 3):
|
||||
image = image.permute(1, 2, 0)
|
||||
image = (image.float().clamp(0, 1) * 255).to(torch.uint8).numpy()
|
||||
return Image.fromarray(image)
|
||||
def to_pixel_values(image_tensor: Tensor) -> Tensor:
|
||||
image = image_tensor.detach().float()
|
||||
if image.shape[-3] == 1:
|
||||
repeats = [1] * image.ndim
|
||||
repeats[-3] = 3
|
||||
image = image.repeat(*repeats)
|
||||
return image
|
||||
|
||||
|
||||
class _FakeVideoEncoder(nn.Module):
|
||||
@@ -242,12 +242,14 @@ class _FakeVideoEncoder(nn.Module):
|
||||
|
||||
|
||||
class _FakeVideoProcessor:
|
||||
def __call__(self, videos, return_tensors: str) -> dict[str, Tensor]:
|
||||
def __call__(self, videos, return_tensors: str, device=None, **kwargs) -> dict[str, Tensor]:
|
||||
assert return_tensors == "pt"
|
||||
if isinstance(videos, list):
|
||||
pixel_values = torch.stack([torch.as_tensor(v) for v in videos])
|
||||
else:
|
||||
pixel_values = torch.as_tensor(videos).unsqueeze(0)
|
||||
if device is not None:
|
||||
pixel_values = pixel_values.to(device)
|
||||
return {"pixel_values_videos": pixel_values}
|
||||
|
||||
|
||||
|
||||
@@ -211,40 +211,42 @@ def test_reset_clears_action_queue(patch_vla_jepa_external_models: None) -> None
|
||||
|
||||
|
||||
def test_prepare_model_inputs_training_format(patch_vla_jepa_external_models: None) -> None:
|
||||
from PIL import Image
|
||||
|
||||
policy = VLAJEPAPolicy(make_config())
|
||||
examples = policy._prepare_model_inputs(make_train_batch())
|
||||
inputs = policy._prepare_model_inputs(make_train_batch())
|
||||
|
||||
assert len(examples) == BATCH_SIZE
|
||||
for ex in examples:
|
||||
assert set(ex) >= {"image", "video", "lang", "action", "state"}
|
||||
assert len(ex["image"]) == 1 and isinstance(ex["image"][0], Image.Image)
|
||||
assert ex["video"].ndim == 5 and ex["video"].dtype == np.uint8 # [V,T,H,W,C]
|
||||
assert ex["action"].shape == (ACTION_HORIZON, ACTION_DIM)
|
||||
assert ex["state"].shape == (1, STATE_DIM)
|
||||
assert set(inputs) >= {"images", "instructions", "videos", "actions", "state"}
|
||||
# images: per-sample, per-view [C, H, W] float tensors (kept as a list for Qwen messages)
|
||||
assert len(inputs["images"]) == BATCH_SIZE and len(inputs["images"][0]) == 1
|
||||
img = inputs["images"][0][0]
|
||||
assert isinstance(img, torch.Tensor) and img.dtype == torch.float32 and img.ndim == 3
|
||||
assert len(inputs["instructions"]) == BATCH_SIZE
|
||||
# videos: batched [B, V, T, C, H, W] float
|
||||
assert inputs["videos"].ndim == 6 and inputs["videos"].shape[0] == BATCH_SIZE
|
||||
assert inputs["videos"].dtype == torch.float32
|
||||
assert inputs["actions"].shape == (BATCH_SIZE, ACTION_HORIZON, ACTION_DIM)
|
||||
assert inputs["state"].shape == (BATCH_SIZE, 1, STATE_DIM)
|
||||
|
||||
|
||||
def test_prepare_model_inputs_inference_omits_action(patch_vla_jepa_external_models: None) -> None:
|
||||
policy = VLAJEPAPolicy(make_config())
|
||||
for ex in policy._prepare_model_inputs(make_inference_batch()):
|
||||
assert "action" not in ex
|
||||
assert "image" in ex and "video" in ex and "lang" in ex
|
||||
inputs = policy._prepare_model_inputs(make_inference_batch())
|
||||
assert "actions" not in inputs and "action_is_pad" not in inputs
|
||||
assert {"images", "instructions", "state"} <= set(inputs)
|
||||
|
||||
|
||||
def test_prepare_model_inputs_missing_task_uses_default(patch_vla_jepa_external_models: None) -> None:
|
||||
policy = VLAJEPAPolicy(make_config())
|
||||
batch = make_inference_batch()
|
||||
del batch["task"]
|
||||
examples = policy._prepare_model_inputs(batch)
|
||||
assert all(isinstance(ex["lang"], str) and len(ex["lang"]) > 0 for ex in examples)
|
||||
instructions = policy._prepare_model_inputs(batch)["instructions"]
|
||||
assert all(isinstance(s, str) and len(s) > 0 for s in instructions)
|
||||
|
||||
|
||||
def test_prepare_model_inputs_string_task_broadcast(patch_vla_jepa_external_models: None) -> None:
|
||||
policy = VLAJEPAPolicy(make_config())
|
||||
batch = make_inference_batch()
|
||||
batch["task"] = "open the drawer"
|
||||
assert all(ex["lang"] == "open the drawer" for ex in policy._prepare_model_inputs(batch))
|
||||
assert policy._prepare_model_inputs(batch)["instructions"] == ["open the drawer"] * BATCH_SIZE
|
||||
|
||||
|
||||
def test_prepare_model_inputs_no_state_omitted(patch_vla_jepa_external_models: None) -> None:
|
||||
@@ -253,7 +255,7 @@ def test_prepare_model_inputs_no_state_omitted(patch_vla_jepa_external_models: N
|
||||
policy = VLAJEPAPolicy(make_config())
|
||||
batch = make_inference_batch()
|
||||
del batch[OBS_STATE]
|
||||
assert all("state" not in ex for ex in policy._prepare_model_inputs(batch))
|
||||
assert "state" not in policy._prepare_model_inputs(batch)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -446,14 +448,14 @@ def test_postprocessor_applied_after_predict_action_chunk(
|
||||
"""
|
||||
from lerobot.policies.vla_jepa.processor_vla_jepa import make_vla_jepa_pre_post_processors
|
||||
|
||||
raw_actions = np.zeros((BATCH_SIZE, ACTION_HORIZON, ACTION_DIM), dtype=np.float32)
|
||||
raw_actions = torch.zeros((BATCH_SIZE, ACTION_HORIZON, ACTION_DIM), dtype=torch.float32)
|
||||
|
||||
cfg = make_config()
|
||||
cfg.clip_normalized_actions = False
|
||||
cfg.binarize_gripper_action = False
|
||||
policy = VLAJEPAPolicy(cfg)
|
||||
policy.eval()
|
||||
monkeypatch.setattr(policy.model, "predict_action", lambda *a, **kw: raw_actions.copy())
|
||||
monkeypatch.setattr(policy.model, "predict_action", lambda *a, **kw: raw_actions.clone())
|
||||
|
||||
dataset_stats = _make_dataset_stats()
|
||||
_, postprocessor = make_vla_jepa_pre_post_processors(cfg, dataset_stats)
|
||||
@@ -564,9 +566,9 @@ def test_single_view_is_duplicated_for_world_model(patch_vla_jepa_external_model
|
||||
original_processor = policy.model.video_processor
|
||||
|
||||
class _CapturingProcessor:
|
||||
def __call__(self, videos: list, return_tensors: str) -> dict:
|
||||
def __call__(self, videos: list, return_tensors: str, **kwargs) -> dict:
|
||||
captured_videos.extend(videos)
|
||||
return original_processor(videos=videos, return_tensors=return_tensors)
|
||||
return original_processor(videos=videos, return_tensors=return_tensors, **kwargs)
|
||||
|
||||
policy.model.video_processor = _CapturingProcessor()
|
||||
policy.forward(_make_multiview_train_batch(num_views=1))
|
||||
@@ -587,9 +589,9 @@ def test_excess_views_trimmed_for_world_model(patch_vla_jepa_external_models: No
|
||||
original_processor = policy.model.video_processor
|
||||
|
||||
class _CapturingProcessor:
|
||||
def __call__(self, videos: list, return_tensors: str) -> dict:
|
||||
def __call__(self, videos: list, return_tensors: str, **kwargs) -> dict:
|
||||
captured_videos.extend(videos)
|
||||
return original_processor(videos=videos, return_tensors=return_tensors)
|
||||
return original_processor(videos=videos, return_tensors=return_tensors, **kwargs)
|
||||
|
||||
policy.model.video_processor = _CapturingProcessor()
|
||||
policy.forward(_make_multiview_train_batch(num_views=3))
|
||||
|
||||
@@ -27,6 +27,7 @@ from lerobot.scripts.lerobot_edit_dataset import (
|
||||
MergeConfig,
|
||||
ModifyTasksConfig,
|
||||
OperationConfig,
|
||||
ReencodeVideosConfig,
|
||||
RemoveFeatureConfig,
|
||||
SplitConfig,
|
||||
_validate_config,
|
||||
@@ -103,3 +104,47 @@ class TestOperationTypeParsing:
|
||||
)
|
||||
resolved_name = OperationConfig.get_choice_name(type(cfg.operation))
|
||||
assert resolved_name == type_name
|
||||
|
||||
|
||||
class TestDepthEncoderParsing:
|
||||
"""Test that the depth encoder is exposed and parsed for video operations."""
|
||||
|
||||
def test_reencode_has_default_depth_encoder(self):
|
||||
cfg = parse_cfg(["--repo_id", "test/repo", "--operation.type", "reencode_videos"])
|
||||
assert isinstance(cfg.operation, ReencodeVideosConfig)
|
||||
# A depth encoder is configured by default so depth videos are re-encoded too.
|
||||
assert cfg.operation.depth_encoder is not None
|
||||
assert hasattr(cfg.operation.depth_encoder, "depth_min")
|
||||
|
||||
def test_reencode_parses_depth_encoder_overrides(self):
|
||||
cfg = parse_cfg(
|
||||
[
|
||||
"--repo_id",
|
||||
"test/repo",
|
||||
"--operation.type",
|
||||
"reencode_videos",
|
||||
"--operation.depth_encoder.extra_options",
|
||||
'{"x265-params": "lossless=1"}',
|
||||
"--operation.depth_encoder.depth_max",
|
||||
"12.0",
|
||||
"--operation.depth_encoder.use_log",
|
||||
"false",
|
||||
]
|
||||
)
|
||||
assert cfg.operation.depth_encoder.extra_options == {"x265-params": "lossless=1"}
|
||||
assert cfg.operation.depth_encoder.depth_max == 12.0
|
||||
assert cfg.operation.depth_encoder.use_log is False
|
||||
|
||||
def test_convert_image_to_video_parses_depth_encoder_overrides(self):
|
||||
cfg = parse_cfg(
|
||||
[
|
||||
"--repo_id",
|
||||
"test/repo",
|
||||
"--operation.type",
|
||||
"convert_image_to_video",
|
||||
"--operation.depth_encoder.depth_min",
|
||||
"0.05",
|
||||
]
|
||||
)
|
||||
assert isinstance(cfg.operation, ConvertImageToVideoConfig)
|
||||
assert cfg.operation.depth_encoder.depth_min == 0.05
|
||||
|
||||
@@ -0,0 +1,67 @@
|
||||
# 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.
|
||||
|
||||
import sys
|
||||
|
||||
import draccus
|
||||
import pytest
|
||||
|
||||
# Importing lerobot_train eagerly pulls in lerobot.datasets, which needs the
|
||||
# `dataset` extra. The base CI tier runs without it, so skip the whole module there.
|
||||
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
||||
|
||||
from lerobot.configs.train import TrainPipelineConfig # noqa: E402
|
||||
from lerobot.policies.act.configuration_act import (
|
||||
ACTConfig, # noqa: E402, F401 (registers --policy.type act)
|
||||
)
|
||||
from lerobot.scripts.lerobot_train import _remote_target_in_argv, train # noqa: E402
|
||||
|
||||
|
||||
def _set_argv(monkeypatch, *args):
|
||||
monkeypatch.setattr(sys, "argv", ["lerobot-train", *args])
|
||||
|
||||
|
||||
def test_remote_target_detected_space_separated(monkeypatch):
|
||||
_set_argv(monkeypatch, "--policy.type", "act", "--job.target", "a10g-small")
|
||||
assert _remote_target_in_argv() is True
|
||||
|
||||
|
||||
def test_remote_target_detected_equals(monkeypatch):
|
||||
_set_argv(monkeypatch, "--job.target=t4-small")
|
||||
assert _remote_target_in_argv() is True
|
||||
|
||||
|
||||
def test_local_string_is_not_remote(monkeypatch):
|
||||
_set_argv(monkeypatch, "--job.target", "local")
|
||||
assert _remote_target_in_argv() is False
|
||||
|
||||
|
||||
def test_no_target_is_not_remote(monkeypatch):
|
||||
_set_argv(monkeypatch, "--policy.type", "act")
|
||||
assert _remote_target_in_argv() is False
|
||||
|
||||
|
||||
def test_train_dispatches_to_submit_when_remote(monkeypatch):
|
||||
"""A remote --job.target short-circuits train() to the HF Jobs submitter."""
|
||||
import lerobot.scripts.lerobot_train as train_module
|
||||
|
||||
captured = []
|
||||
monkeypatch.setattr(train_module, "submit_to_hf", lambda cfg: captured.append(cfg) or "submitted")
|
||||
cfg = draccus.parse(
|
||||
TrainPipelineConfig,
|
||||
args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"],
|
||||
)
|
||||
# Returns the submitter's result and never enters the local training path.
|
||||
assert train(cfg) == "submitted"
|
||||
assert captured == [cfg]
|
||||
@@ -58,7 +58,46 @@ def download_dataset(repo_id, episodes):
|
||||
print(f"Dataset {repo_id} downloaded successfully")
|
||||
|
||||
|
||||
def run_accelerate_training(config_args, num_processes=4, temp_dir=None):
|
||||
def _write_multi_gpu_config(f, num_processes):
|
||||
f.write("compute_environment: LOCAL_MACHINE\n")
|
||||
f.write("distributed_type: MULTI_GPU\n")
|
||||
f.write("mixed_precision: 'no'\n")
|
||||
f.write(f"num_processes: {num_processes}\n")
|
||||
f.write("use_cpu: false\n")
|
||||
f.write("gpu_ids: all\n")
|
||||
f.write("downcast_bf16: 'no'\n")
|
||||
f.write("machine_rank: 0\n")
|
||||
f.write("main_training_function: main\n")
|
||||
f.write("num_machines: 1\n")
|
||||
f.write("rdzv_backend: static\n")
|
||||
f.write("same_network: true\n")
|
||||
|
||||
|
||||
def _write_fsdp_config(f, num_processes):
|
||||
# FSDP1 with FULL_SHARD (ZeRO-3-equivalent) and FULL_STATE_DICT, matching
|
||||
# docs/source/multi_gpu_training.mdx. ACT's repeated transformer blocks are the wrap units;
|
||||
# fsdp_use_orig_params is required because LeRobot builds the optimizer before prepare().
|
||||
f.write("compute_environment: LOCAL_MACHINE\n")
|
||||
f.write("distributed_type: FSDP\n")
|
||||
f.write("mixed_precision: 'no'\n")
|
||||
f.write(f"num_processes: {num_processes}\n")
|
||||
f.write("use_cpu: false\n")
|
||||
f.write("gpu_ids: all\n")
|
||||
f.write("machine_rank: 0\n")
|
||||
f.write("main_training_function: main\n")
|
||||
f.write("num_machines: 1\n")
|
||||
f.write("rdzv_backend: static\n")
|
||||
f.write("same_network: true\n")
|
||||
f.write("fsdp_config:\n")
|
||||
f.write(" fsdp_version: 1\n")
|
||||
f.write(" fsdp_sharding_strategy: FULL_SHARD\n")
|
||||
f.write(" fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP\n")
|
||||
f.write(" fsdp_transformer_layer_cls_to_wrap: ACTEncoderLayer,ACTDecoderLayer\n")
|
||||
f.write(" fsdp_use_orig_params: true\n")
|
||||
f.write(" fsdp_state_dict_type: FULL_STATE_DICT\n")
|
||||
|
||||
|
||||
def run_accelerate_training(config_args, num_processes=4, temp_dir=None, distributed_type="MULTI_GPU"):
|
||||
"""
|
||||
Helper function to run training with accelerate launch.
|
||||
|
||||
@@ -66,6 +105,7 @@ def run_accelerate_training(config_args, num_processes=4, temp_dir=None):
|
||||
config_args: List of config arguments to pass to lerobot_train.py
|
||||
num_processes: Number of processes (GPUs) to use
|
||||
temp_dir: Temporary directory for outputs
|
||||
distributed_type: "MULTI_GPU" (DDP) or "FSDP" — selects the generated accelerate config.
|
||||
|
||||
Returns:
|
||||
subprocess.CompletedProcess result
|
||||
@@ -75,18 +115,10 @@ def run_accelerate_training(config_args, num_processes=4, temp_dir=None):
|
||||
|
||||
# Write YAML config
|
||||
with open(config_path, "w") as f:
|
||||
f.write("compute_environment: LOCAL_MACHINE\n")
|
||||
f.write("distributed_type: MULTI_GPU\n")
|
||||
f.write("mixed_precision: 'no'\n")
|
||||
f.write(f"num_processes: {num_processes}\n")
|
||||
f.write("use_cpu: false\n")
|
||||
f.write("gpu_ids: all\n")
|
||||
f.write("downcast_bf16: 'no'\n")
|
||||
f.write("machine_rank: 0\n")
|
||||
f.write("main_training_function: main\n")
|
||||
f.write("num_machines: 1\n")
|
||||
f.write("rdzv_backend: static\n")
|
||||
f.write("same_network: true\n")
|
||||
if distributed_type == "FSDP":
|
||||
_write_fsdp_config(f, num_processes)
|
||||
else:
|
||||
_write_multi_gpu_config(f, num_processes)
|
||||
|
||||
cmd = [
|
||||
"accelerate",
|
||||
@@ -134,7 +166,7 @@ class TestMultiGPUTraining:
|
||||
f"--output_dir={output_dir}",
|
||||
"--batch_size=4",
|
||||
"--steps=10",
|
||||
"--eval_freq=-1",
|
||||
"--env_eval_freq=-1",
|
||||
"--log_freq=5",
|
||||
"--save_freq=10",
|
||||
"--seed=42",
|
||||
@@ -177,7 +209,7 @@ class TestMultiGPUTraining:
|
||||
f"--output_dir={output_dir}",
|
||||
"--batch_size=4",
|
||||
"--steps=20",
|
||||
"--eval_freq=-1",
|
||||
"--env_eval_freq=-1",
|
||||
"--log_freq=5",
|
||||
"--save_freq=10",
|
||||
"--seed=42",
|
||||
@@ -211,3 +243,66 @@ class TestMultiGPUTraining:
|
||||
# Verify optimizer state exists
|
||||
optimizer_state = training_state_dir / "optimizer_state.safetensors"
|
||||
assert optimizer_state.exists(), f"No optimizer state in checkpoint {checkpoint_dir}"
|
||||
|
||||
def test_fsdp_optimizer_save_and_resume(self):
|
||||
"""
|
||||
Test that FSDP saves the (gathered) optimizer state and can resume from it.
|
||||
|
||||
Trains a few steps under FSDP, verifies the gathered optimizer state is written next to the
|
||||
rest of the training state, then resumes from the checkpoint for more steps and checks it
|
||||
completes without shape/key errors in the FSDP optimizer load path.
|
||||
"""
|
||||
# Pre-download dataset to avoid race conditions
|
||||
download_dataset("lerobot/pusht", episodes=[0])
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
output_dir = Path(temp_dir) / "outputs"
|
||||
|
||||
config_args = [
|
||||
"--dataset.repo_id=lerobot/pusht",
|
||||
"--dataset.episodes=[0]",
|
||||
"--policy.type=act",
|
||||
"--policy.device=cuda",
|
||||
"--policy.push_to_hub=false",
|
||||
f"--output_dir={output_dir}",
|
||||
"--batch_size=4",
|
||||
"--steps=10",
|
||||
"--env_eval_freq=-1",
|
||||
"--log_freq=5",
|
||||
"--save_freq=10",
|
||||
"--seed=42",
|
||||
"--num_workers=0",
|
||||
]
|
||||
|
||||
result = run_accelerate_training(
|
||||
config_args, num_processes=2, temp_dir=temp_dir, distributed_type="FSDP"
|
||||
)
|
||||
assert result.returncode == 0, (
|
||||
f"FSDP training failed:\nSTDOUT:\n{result.stdout}\n\nSTDERR:\n{result.stderr}"
|
||||
)
|
||||
|
||||
# The gathered optimizer state must be written under FSDP (proves the save collective ran),
|
||||
# in the same safetensors format as single-GPU training.
|
||||
training_state_dir = output_dir / "checkpoints" / "last" / "training_state"
|
||||
optimizer_state = training_state_dir / "optimizer_state.safetensors"
|
||||
optimizer_param_groups = training_state_dir / "optimizer_param_groups.json"
|
||||
assert optimizer_state.exists(), f"FSDP optimizer state not saved in {training_state_dir}"
|
||||
assert optimizer_param_groups.exists(), (
|
||||
f"FSDP optimizer param groups not saved in {training_state_dir}"
|
||||
)
|
||||
|
||||
# Resume from the checkpoint for more steps. A successful run proves load_fsdp_optimizer
|
||||
# accepts the saved state and reshards it without shape/key errors.
|
||||
resume_config = output_dir / "checkpoints" / "last" / "pretrained_model" / "train_config.json"
|
||||
resume_args = [
|
||||
f"--config_path={resume_config}",
|
||||
"--resume=true",
|
||||
"--steps=20",
|
||||
]
|
||||
resume_result = run_accelerate_training(
|
||||
resume_args, num_processes=2, temp_dir=temp_dir, distributed_type="FSDP"
|
||||
)
|
||||
assert resume_result.returncode == 0, (
|
||||
f"FSDP resume failed:\nSTDOUT:\n{resume_result.stdout}\n\nSTDERR:\n{resume_result.stderr}"
|
||||
)
|
||||
assert "End of training" in resume_result.stdout or "End of training" in resume_result.stderr
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
# 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.
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from lerobot.utils.hub import find_latest_hub_checkpoint
|
||||
|
||||
|
||||
def _patch_list_files(monkeypatch, files):
|
||||
api = MagicMock()
|
||||
api.list_repo_files.return_value = files
|
||||
# HfApi is imported into lerobot.utils.hub at module load, so patch it there.
|
||||
monkeypatch.setattr("lerobot.utils.hub.HfApi", lambda *a, **k: api)
|
||||
return api
|
||||
|
||||
|
||||
def test_find_latest_hub_checkpoint_picks_highest_step(monkeypatch):
|
||||
_patch_list_files(
|
||||
monkeypatch,
|
||||
[
|
||||
"README.md",
|
||||
"checkpoints/000500/pretrained_model/model.safetensors",
|
||||
"checkpoints/000500/training_state/training_step.json",
|
||||
"checkpoints/020000/pretrained_model/model.safetensors",
|
||||
"checkpoints/001000/training_state/training_step.json",
|
||||
],
|
||||
)
|
||||
# Numeric max, not lexicographic — "020000" beats "001000"/"000500".
|
||||
assert find_latest_hub_checkpoint("u/run") == "checkpoints/020000"
|
||||
|
||||
|
||||
def test_find_latest_hub_checkpoint_ignores_non_step_entries(monkeypatch):
|
||||
_patch_list_files(
|
||||
monkeypatch,
|
||||
["checkpoints/last/pretrained_model/model.safetensors", "config.json"],
|
||||
)
|
||||
# "last" (a symlink target name) is not a numeric step → no resolvable checkpoint.
|
||||
assert find_latest_hub_checkpoint("u/run") is None
|
||||
|
||||
|
||||
def test_find_latest_hub_checkpoint_none_when_no_checkpoints(monkeypatch):
|
||||
_patch_list_files(monkeypatch, ["config.json", "model.safetensors"])
|
||||
assert find_latest_hub_checkpoint("u/run") is None
|
||||
@@ -0,0 +1,228 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 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.
|
||||
|
||||
"""Unit tests for the display-independent keyboard input helpers.
|
||||
|
||||
These cover the parts most likely to regress: the environment-detection decision
|
||||
table (the heart of the Wayland/headless fix), the macOS trust probe, the control
|
||||
mapping, the terminal escape-sequence parsing, and backend selection. They require
|
||||
neither ``pynput`` nor a real terminal.
|
||||
"""
|
||||
|
||||
import io
|
||||
import platform
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
import lerobot.utils.keyboard_input as ki
|
||||
from lerobot.utils.keyboard_input import (
|
||||
TerminalKeyListener,
|
||||
apply_recording_control,
|
||||
create_key_listener,
|
||||
init_keyboard_listener,
|
||||
is_headless,
|
||||
is_wayland,
|
||||
pynput_can_capture,
|
||||
pynput_listener_is_trusted,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_detection_caches():
|
||||
"""The detection helpers are ``@cache``-decorated; clear around each test."""
|
||||
for fn in (is_headless, is_wayland, pynput_can_capture):
|
||||
fn.cache_clear()
|
||||
yield
|
||||
for fn in (is_headless, is_wayland, pynput_can_capture):
|
||||
fn.cache_clear()
|
||||
|
||||
|
||||
def _set_platform(monkeypatch, name):
|
||||
monkeypatch.setattr(platform, "system", lambda: name)
|
||||
|
||||
|
||||
def _set_tty(monkeypatch, is_tty):
|
||||
stdin = io.StringIO("")
|
||||
stdin.isatty = lambda: is_tty
|
||||
monkeypatch.setattr(sys, "stdin", stdin)
|
||||
|
||||
|
||||
# --- Environment detection (the core of the fix) ---------------------------
|
||||
@pytest.mark.parametrize(
|
||||
("system", "env", "expected"),
|
||||
[
|
||||
("Linux", {}, True), # no display server
|
||||
("Linux", {"DISPLAY": ":0"}, False), # X11
|
||||
("Linux", {"WAYLAND_DISPLAY": "wayland-0"}, False), # Wayland
|
||||
("Darwin", {}, False), # display always assumed present
|
||||
],
|
||||
)
|
||||
def test_is_headless(monkeypatch, system, env, expected):
|
||||
_set_platform(monkeypatch, system)
|
||||
monkeypatch.delenv("DISPLAY", raising=False)
|
||||
monkeypatch.delenv("WAYLAND_DISPLAY", raising=False)
|
||||
for key, value in env.items():
|
||||
monkeypatch.setenv(key, value)
|
||||
assert is_headless() is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("env", "expected"),
|
||||
[
|
||||
({"XDG_SESSION_TYPE": "wayland"}, True),
|
||||
({"WAYLAND_DISPLAY": "wayland-0"}, True),
|
||||
({"XDG_SESSION_TYPE": "x11"}, False),
|
||||
({}, False),
|
||||
],
|
||||
)
|
||||
def test_is_wayland(monkeypatch, env, expected):
|
||||
monkeypatch.delenv("XDG_SESSION_TYPE", raising=False)
|
||||
monkeypatch.delenv("WAYLAND_DISPLAY", raising=False)
|
||||
for key, value in env.items():
|
||||
monkeypatch.setenv(key, value)
|
||||
assert is_wayland() is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("system", "env", "pynput_available", "expected"),
|
||||
[
|
||||
("Linux", {"DISPLAY": ":0"}, True, True), # X11
|
||||
("Linux", {"DISPLAY": ":0", "WAYLAND_DISPLAY": "wayland-0"}, True, False), # Wayland
|
||||
("Linux", {}, True, False), # headless
|
||||
("Darwin", {}, True, True),
|
||||
("Linux", {"DISPLAY": ":0"}, False, False), # pynput not installed
|
||||
],
|
||||
)
|
||||
def test_pynput_can_capture(monkeypatch, system, env, pynput_available, expected):
|
||||
_set_platform(monkeypatch, system)
|
||||
monkeypatch.setattr(ki, "_pynput_available", pynput_available)
|
||||
for var in ("DISPLAY", "WAYLAND_DISPLAY", "XDG_SESSION_TYPE"):
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
for key, value in env.items():
|
||||
monkeypatch.setenv(key, value)
|
||||
assert pynput_can_capture() is expected
|
||||
|
||||
|
||||
# --- macOS trust probe ------------------------------------------------------
|
||||
class _FakeListener:
|
||||
def __init__(self, is_trusted):
|
||||
self.IS_TRUSTED = is_trusted
|
||||
|
||||
|
||||
def test_pynput_listener_is_trusted(monkeypatch):
|
||||
_set_platform(monkeypatch, "Linux")
|
||||
assert pynput_listener_is_trusted(_FakeListener(False)) is True # non-macOS: always assumed ok
|
||||
_set_platform(monkeypatch, "Darwin")
|
||||
assert pynput_listener_is_trusted(_FakeListener(False), timeout_s=0.05) is False
|
||||
|
||||
|
||||
# --- Control mapping --------------------------------------------------------
|
||||
def test_apply_recording_control():
|
||||
events = {"exit_early": False, "rerecord_episode": False, "stop_recording": False}
|
||||
apply_recording_control("left", events)
|
||||
assert events == {"exit_early": True, "rerecord_episode": True, "stop_recording": False}
|
||||
apply_recording_control("esc", events)
|
||||
assert events["stop_recording"] is True
|
||||
apply_recording_control("up", events) # unknown control -> no-op (no error)
|
||||
|
||||
|
||||
# --- Terminal escape-sequence parsing (the tricky bit) ----------------------
|
||||
def _drive(listener, byte_seq):
|
||||
"""Run the listener's read loop over a scripted list of bytes (no real terminal)."""
|
||||
script = list(byte_seq)
|
||||
|
||||
def fake_read(timeout):
|
||||
if script:
|
||||
return script.pop(0)
|
||||
listener._running = False
|
||||
return None
|
||||
|
||||
listener._read_char = fake_read
|
||||
listener._running = True
|
||||
listener._run()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("byte_seq", "expected"),
|
||||
[
|
||||
(["\x1b", "[", "C"], ["right"]), # CSI arrow
|
||||
(["\x1b", "O", "D"], ["left"]), # SS3 arrow (e.g. over SSH/tmux)
|
||||
(["\x1b"], ["esc"]), # bare ESC
|
||||
(["\x1b", "[", "A"], ["up"]), # decoded even though the record handler ignores it
|
||||
(["n"], ["n"]), # letter passthrough
|
||||
],
|
||||
)
|
||||
def test_terminal_parsing(byte_seq, expected):
|
||||
collected = []
|
||||
_drive(TerminalKeyListener(collected.append), byte_seq)
|
||||
assert collected == expected
|
||||
|
||||
|
||||
# --- Backend selection ------------------------------------------------------
|
||||
def test_init_selects_terminal_when_pynput_cannot_capture(monkeypatch):
|
||||
monkeypatch.setattr(ki, "pynput_can_capture", lambda: False)
|
||||
_set_tty(monkeypatch, is_tty=True)
|
||||
monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None) # avoid touching termios
|
||||
listener, _ = init_keyboard_listener()
|
||||
assert isinstance(listener, TerminalKeyListener)
|
||||
|
||||
|
||||
def test_init_returns_none_without_tty(monkeypatch):
|
||||
monkeypatch.setattr(ki, "pynput_can_capture", lambda: False)
|
||||
_set_tty(monkeypatch, is_tty=False)
|
||||
listener, _ = init_keyboard_listener()
|
||||
assert listener is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("key", "flag"),
|
||||
[("right", "exit_early"), ("r", "rerecord_episode"), ("q", "stop_recording")],
|
||||
)
|
||||
def test_init_terminal_key_routing(monkeypatch, key, flag):
|
||||
"""Arrows and their letter equivalents drive the same events (terminal backend)."""
|
||||
monkeypatch.setattr(ki, "pynput_can_capture", lambda: False)
|
||||
_set_tty(monkeypatch, is_tty=True)
|
||||
monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None)
|
||||
listener, events = init_keyboard_listener()
|
||||
listener._on_key(key)
|
||||
assert events[flag] is True
|
||||
|
||||
|
||||
# --- Shared factory + pynput key resolver -----------------------------------
|
||||
def test_resolve_pynput_key_char_fallback():
|
||||
"""Unmapped keys fall back to ``.char`` (and yield None when there is none)."""
|
||||
assert ki._resolve_pynput_key(type("K", (), {"char": "s"})()) == "s"
|
||||
assert ki._resolve_pynput_key(type("K", (), {"char": None})()) is None
|
||||
assert ki._resolve_pynput_key(type("K", (), {"char": ""})()) is None # empty char -> no key
|
||||
|
||||
|
||||
def test_create_key_listener_routes_to_dispatch(monkeypatch):
|
||||
"""The terminal backend forwards canonical key names straight to ``dispatch``."""
|
||||
monkeypatch.setattr(ki, "pynput_can_capture", lambda: False)
|
||||
_set_tty(monkeypatch, is_tty=True)
|
||||
monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None)
|
||||
seen = []
|
||||
listener = create_key_listener(seen.append, controls_help="save='s'")
|
||||
assert isinstance(listener, TerminalKeyListener)
|
||||
listener._on_key("space")
|
||||
assert seen == ["space"]
|
||||
|
||||
|
||||
def test_create_key_listener_none_without_tty(monkeypatch):
|
||||
monkeypatch.setattr(ki, "pynput_can_capture", lambda: False)
|
||||
_set_tty(monkeypatch, is_tty=False)
|
||||
assert create_key_listener(lambda name: None) is None
|
||||
@@ -15,7 +15,9 @@
|
||||
# limitations under the License.
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import Mock, patch
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from lerobot.common.train_utils import (
|
||||
get_step_checkpoint_dir,
|
||||
@@ -24,6 +26,7 @@ from lerobot.common.train_utils import (
|
||||
load_training_num_processes,
|
||||
load_training_state,
|
||||
load_training_step,
|
||||
push_checkpoint_to_hub,
|
||||
save_checkpoint,
|
||||
save_training_state,
|
||||
save_training_step,
|
||||
@@ -136,3 +139,87 @@ def test_save_load_training_state(tmp_path, optimizer, scheduler):
|
||||
assert loaded_step == 10
|
||||
assert loaded_optimizer is optimizer
|
||||
assert loaded_scheduler is scheduler
|
||||
|
||||
|
||||
def test_load_training_state_skip_optimizer(tmp_path, optimizer, scheduler):
|
||||
# FSDP loads optimizer separately (after accelerator.prepare)
|
||||
# load_training_state(load_optimizer=False) must restore step + scheduler but leave the
|
||||
# optimizer untouched and never touch the on-disk optimizer state.
|
||||
save_training_state(tmp_path, 10, optimizer, scheduler)
|
||||
with patch("lerobot.common.train_utils.load_optimizer_state") as mock_load_optimizer_state:
|
||||
loaded_step, loaded_optimizer, loaded_scheduler = load_training_state(
|
||||
tmp_path, optimizer, scheduler, load_optimizer=False
|
||||
)
|
||||
mock_load_optimizer_state.assert_not_called()
|
||||
assert loaded_step == 10
|
||||
assert loaded_optimizer is optimizer
|
||||
assert loaded_scheduler is scheduler
|
||||
|
||||
|
||||
def test_push_checkpoint_to_hub_creates_repo_and_uploads(tmp_path, monkeypatch):
|
||||
ckpt = tmp_path / "010000"
|
||||
(ckpt / "pretrained_model").mkdir(parents=True)
|
||||
api = MagicMock()
|
||||
monkeypatch.setattr("lerobot.common.train_utils.HfApi", lambda *a, **k: api)
|
||||
push_checkpoint_to_hub(ckpt, "user/run", private=True)
|
||||
api.create_repo.assert_called_once()
|
||||
assert api.create_repo.call_args.kwargs["private"] is True
|
||||
assert api.create_repo.call_args.kwargs["repo_type"] == "model"
|
||||
api.upload_folder.assert_called_once()
|
||||
kwargs = api.upload_folder.call_args.kwargs
|
||||
assert kwargs["repo_id"] == "user/run"
|
||||
assert kwargs["repo_type"] == "model"
|
||||
assert kwargs["path_in_repo"] == "checkpoints/010000"
|
||||
assert kwargs["folder_path"] == str(ckpt)
|
||||
assert kwargs["commit_message"] == "checkpoint 010000"
|
||||
# A tag named after the checkpoint step is created so the checkpoint can be
|
||||
# recovered with --policy.pretrained_revision instead of a commit sha.
|
||||
api.create_tag.assert_called_once()
|
||||
tag_kwargs = api.create_tag.call_args.kwargs
|
||||
assert tag_kwargs["tag"] == "010000"
|
||||
assert tag_kwargs["revision"] == api.upload_folder.return_value.oid
|
||||
assert tag_kwargs["repo_type"] == "model"
|
||||
assert tag_kwargs["exist_ok"] is True
|
||||
|
||||
|
||||
def test_push_checkpoint_to_hub_defaults_to_hub_default_visibility(tmp_path, monkeypatch):
|
||||
ckpt = tmp_path / "010000"
|
||||
(ckpt / "pretrained_model").mkdir(parents=True)
|
||||
api = MagicMock()
|
||||
monkeypatch.setattr("lerobot.common.train_utils.HfApi", lambda *a, **k: api)
|
||||
push_checkpoint_to_hub(ckpt, "user/run")
|
||||
api.create_repo.assert_called_once()
|
||||
assert api.create_repo.call_args.kwargs["private"] is None
|
||||
|
||||
|
||||
def test_resolve_resume_checkpoint_downloads_latest_and_links(tmp_path, monkeypatch):
|
||||
from lerobot.common import train_utils
|
||||
|
||||
out = tmp_path / "run"
|
||||
|
||||
def fake_snapshot_download(repo_id, repo_type, allow_patterns, local_dir):
|
||||
# Mimic the Hub layout the real download materializes locally.
|
||||
assert allow_patterns == "checkpoints/020000/*"
|
||||
(Path(local_dir) / "checkpoints" / "020000" / "pretrained_model").mkdir(parents=True)
|
||||
return local_dir
|
||||
|
||||
monkeypatch.setattr("lerobot.common.train_utils.snapshot_download", fake_snapshot_download)
|
||||
monkeypatch.setattr(
|
||||
"lerobot.common.train_utils.find_latest_hub_checkpoint", lambda repo_id: "checkpoints/020000"
|
||||
)
|
||||
|
||||
checkpoint_dir = train_utils.resolve_resume_checkpoint("u/run", out)
|
||||
|
||||
assert checkpoint_dir == out / CHECKPOINTS_DIR / "020000"
|
||||
last = out / CHECKPOINTS_DIR / LAST_CHECKPOINT_LINK
|
||||
assert last.is_symlink()
|
||||
# `last` points at the downloaded step dir.
|
||||
assert (last.parent / last.readlink()).resolve() == checkpoint_dir.resolve()
|
||||
|
||||
|
||||
def test_resolve_resume_checkpoint_raises_without_checkpoints(tmp_path, monkeypatch):
|
||||
from lerobot.common import train_utils
|
||||
|
||||
monkeypatch.setattr("lerobot.common.train_utils.find_latest_hub_checkpoint", lambda repo_id: None)
|
||||
with pytest.raises(FileNotFoundError, match="No checkpoint"):
|
||||
train_utils.resolve_resume_checkpoint("u/run", tmp_path / "run")
|
||||
|
||||
@@ -30,46 +30,77 @@ from lerobot.utils.constants import OBS_STATE
|
||||
@pytest.fixture
|
||||
def mock_rerun(monkeypatch):
|
||||
"""
|
||||
Provide a mock `rerun` module so tests don't depend on the real library.
|
||||
Also reload the module-under-test so it binds to this mock `rr`.
|
||||
Provide a mock `rerun` module (and `rerun.blueprint` submodule) so tests don't
|
||||
depend on the real library. Also reload the module-under-test so it binds to
|
||||
this mock `rr`.
|
||||
"""
|
||||
calls = []
|
||||
blueprints = []
|
||||
|
||||
class DummyScalar:
|
||||
def __init__(self, value):
|
||||
self.value = float(value)
|
||||
# Scalars may be built from a single float or from a 1D array batch.
|
||||
self.value = value
|
||||
|
||||
class DummyImage:
|
||||
def __init__(self, arr):
|
||||
self.arr = arr
|
||||
|
||||
def compress(self, *a, **k):
|
||||
return self
|
||||
|
||||
class DummyDepthImage:
|
||||
def __init__(self, arr, colormap=None):
|
||||
self.arr = arr
|
||||
self.colormap = colormap
|
||||
|
||||
def dummy_log(key, obj=None, **kwargs):
|
||||
# Accept either positional `obj` or keyword `entity` and record remaining kwargs.
|
||||
if obj is None and "entity" in kwargs:
|
||||
obj = kwargs.pop("entity")
|
||||
calls.append((key, obj, kwargs))
|
||||
|
||||
def dummy_send_blueprint(blueprint, *a, **k):
|
||||
blueprints.append(blueprint)
|
||||
|
||||
# Mock the `rerun.blueprint` submodule used to build the layout.
|
||||
dummy_rrb = SimpleNamespace(
|
||||
Spatial2DView=lambda origin=None, name=None: SimpleNamespace(
|
||||
kind="Spatial2DView", origin=origin, name=name
|
||||
),
|
||||
TimeSeriesView=lambda name=None, contents=None: SimpleNamespace(
|
||||
kind="TimeSeriesView", name=name, contents=contents
|
||||
),
|
||||
Grid=lambda *views: SimpleNamespace(kind="Grid", views=list(views)),
|
||||
Blueprint=lambda root: SimpleNamespace(kind="Blueprint", root=root),
|
||||
)
|
||||
|
||||
dummy_rr = SimpleNamespace(
|
||||
__name__="rerun",
|
||||
__package__="rerun",
|
||||
__spec__=SimpleNamespace(name="rerun", submodule_search_locations=None),
|
||||
Scalars=DummyScalar,
|
||||
Image=DummyImage,
|
||||
DepthImage=DummyDepthImage,
|
||||
components=SimpleNamespace(Colormap=SimpleNamespace(Viridis="viridis")),
|
||||
log=dummy_log,
|
||||
send_blueprint=dummy_send_blueprint,
|
||||
init=lambda *a, **k: None,
|
||||
spawn=lambda *a, **k: None,
|
||||
blueprint=dummy_rrb,
|
||||
)
|
||||
|
||||
# Inject fake module into sys.modules
|
||||
# Inject fake modules into sys.modules (both `rerun` and `rerun.blueprint`).
|
||||
monkeypatch.setitem(sys.modules, "rerun", dummy_rr)
|
||||
monkeypatch.setitem(sys.modules, "rerun.blueprint", dummy_rrb)
|
||||
|
||||
# Now import and reload the module under test, to bind to our rerun mock
|
||||
import lerobot.utils.visualization_utils as vu
|
||||
|
||||
importlib.reload(vu)
|
||||
|
||||
# Expose both the reloaded module and the call recorder
|
||||
yield vu, calls
|
||||
# Expose the reloaded module, the call recorder and the captured blueprints
|
||||
yield vu, calls, blueprints
|
||||
|
||||
|
||||
def _keys(calls):
|
||||
@@ -92,8 +123,13 @@ def _kwargs_for(calls, key):
|
||||
raise KeyError(f"Key {key} not found in calls: {calls}")
|
||||
|
||||
|
||||
def _views_by_kind(blueprint, kind):
|
||||
"""Return the views of a given kind from the (single) blueprint's grid."""
|
||||
return [v for v in blueprint.root.views if v.kind == kind]
|
||||
|
||||
|
||||
def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
|
||||
vu, calls = mock_rerun
|
||||
vu, calls, blueprints = mock_rerun
|
||||
|
||||
# Build EnvTransition dict
|
||||
obs = {
|
||||
@@ -103,7 +139,7 @@ def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
|
||||
}
|
||||
act = {
|
||||
"action.throttle": 0.7,
|
||||
# 1D array should log individual Scalars with suffix _i
|
||||
# 1D array should be logged as a single Scalars batch under one entity path
|
||||
"action.vector": np.array([1.0, 2.0], dtype=np.float32),
|
||||
}
|
||||
transition = {
|
||||
@@ -120,31 +156,28 @@ def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
|
||||
# - observation.state.temperature -> Scalars
|
||||
# - observation.camera -> Image (HWC) with static=True
|
||||
# - action.throttle -> Scalars
|
||||
# - action.vector_0, action.vector_1 -> Scalars
|
||||
# - action.vector -> single Scalars batch (no per-element suffix)
|
||||
expected_keys = {
|
||||
f"{OBS_STATE}.temperature",
|
||||
"observation.camera",
|
||||
"action.throttle",
|
||||
"action.vector_0",
|
||||
"action.vector_1",
|
||||
"action.vector",
|
||||
}
|
||||
assert set(_keys(calls)) == expected_keys
|
||||
|
||||
# Check scalar types and values
|
||||
temp_obj = _obj_for(calls, f"{OBS_STATE}.temperature")
|
||||
assert type(temp_obj).__name__ == "DummyScalar"
|
||||
assert temp_obj.value == pytest.approx(25.0)
|
||||
assert float(temp_obj.value) == pytest.approx(25.0)
|
||||
|
||||
throttle_obj = _obj_for(calls, "action.throttle")
|
||||
assert type(throttle_obj).__name__ == "DummyScalar"
|
||||
assert throttle_obj.value == pytest.approx(0.7)
|
||||
assert float(throttle_obj.value) == pytest.approx(0.7)
|
||||
|
||||
v0 = _obj_for(calls, "action.vector_0")
|
||||
v1 = _obj_for(calls, "action.vector_1")
|
||||
assert type(v0).__name__ == "DummyScalar"
|
||||
assert type(v1).__name__ == "DummyScalar"
|
||||
assert v0.value == pytest.approx(1.0)
|
||||
assert v1.value == pytest.approx(2.0)
|
||||
# 1D vector logged as a single batched Scalars under one entity path
|
||||
vec = _obj_for(calls, "action.vector")
|
||||
assert type(vec).__name__ == "DummyScalar"
|
||||
np.testing.assert_allclose(np.asarray(vec.value), [1.0, 2.0])
|
||||
|
||||
# Check image handling: CHW -> HWC
|
||||
img_obj = _obj_for(calls, "observation.camera")
|
||||
@@ -152,9 +185,24 @@ def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
|
||||
assert img_obj.arr.shape == (10, 20, 3) # transposed
|
||||
assert _kwargs_for(calls, "observation.camera").get("static", False) is True # static=True for images
|
||||
|
||||
# A blueprint should have been built and sent exactly once, and cached on the function.
|
||||
assert len(blueprints) == 1
|
||||
assert vu.log_rerun_data.blueprint is blueprints[0]
|
||||
|
||||
bp = blueprints[0]
|
||||
# One spatial view per image path
|
||||
spatial_views = _views_by_kind(bp, "Spatial2DView")
|
||||
assert {v.origin for v in spatial_views} == {"observation.camera"}
|
||||
|
||||
# One time-series view each for observation and action scalars
|
||||
ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
|
||||
assert set(ts_views) == {"observation", "action"}
|
||||
assert ts_views["observation"].contents == [f"{OBS_STATE}.temperature"]
|
||||
assert ts_views["action"].contents == ["action.throttle", "action.vector"]
|
||||
|
||||
|
||||
def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
|
||||
vu, calls = mock_rerun
|
||||
vu, calls, blueprints = mock_rerun
|
||||
|
||||
# First dict without prefixes treated as observation
|
||||
# Second dict without prefixes treated as action
|
||||
@@ -173,14 +221,12 @@ def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
|
||||
# First dict was treated as observation, second as action
|
||||
vu.log_rerun_data(observation=obs_plain, action=act_plain)
|
||||
|
||||
# Expected keys with auto-prefixes
|
||||
# Expected keys with auto-prefixes. The 1D vector is a single batched Scalars.
|
||||
expected = {
|
||||
"observation.temp",
|
||||
"observation.img",
|
||||
"action.throttle",
|
||||
"action.vec_0",
|
||||
"action.vec_1",
|
||||
"action.vec_2",
|
||||
"action.vec",
|
||||
}
|
||||
logged = set(_keys(calls))
|
||||
assert logged == expected
|
||||
@@ -188,11 +234,11 @@ def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
|
||||
# Scalars
|
||||
t = _obj_for(calls, "observation.temp")
|
||||
assert type(t).__name__ == "DummyScalar"
|
||||
assert t.value == pytest.approx(1.5)
|
||||
assert float(t.value) == pytest.approx(1.5)
|
||||
|
||||
throttle = _obj_for(calls, "action.throttle")
|
||||
assert type(throttle).__name__ == "DummyScalar"
|
||||
assert throttle.value == pytest.approx(0.3)
|
||||
assert float(throttle.value) == pytest.approx(0.3)
|
||||
|
||||
# Image stays HWC
|
||||
img = _obj_for(calls, "observation.img")
|
||||
@@ -200,15 +246,23 @@ def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
|
||||
assert img.arr.shape == (5, 6, 3)
|
||||
assert _kwargs_for(calls, "observation.img").get("static", False) is True
|
||||
|
||||
# Vectors
|
||||
for i, val in enumerate([9, 8, 7]):
|
||||
o = _obj_for(calls, f"action.vec_{i}")
|
||||
assert type(o).__name__ == "DummyScalar"
|
||||
assert o.value == pytest.approx(val)
|
||||
# Vector logged as a single batched Scalars under one entity path
|
||||
vec = _obj_for(calls, "action.vec")
|
||||
assert type(vec).__name__ == "DummyScalar"
|
||||
np.testing.assert_allclose(np.asarray(vec.value), [9, 8, 7])
|
||||
|
||||
# Blueprint sent once with the expected view layout
|
||||
assert len(blueprints) == 1
|
||||
bp = blueprints[0]
|
||||
spatial_views = _views_by_kind(bp, "Spatial2DView")
|
||||
assert {v.origin for v in spatial_views} == {"observation.img"}
|
||||
ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
|
||||
assert ts_views["observation"].contents == ["observation.temp"]
|
||||
assert ts_views["action"].contents == ["action.throttle", "action.vec"]
|
||||
|
||||
|
||||
def test_log_rerun_data_kwargs_only(mock_rerun):
|
||||
vu, calls = mock_rerun
|
||||
vu, calls, blueprints = mock_rerun
|
||||
|
||||
vu.log_rerun_data(
|
||||
observation={"observation.temp": 10.0, "observation.gray": np.zeros((8, 8, 1), dtype=np.uint8)},
|
||||
@@ -222,13 +276,35 @@ def test_log_rerun_data_kwargs_only(mock_rerun):
|
||||
|
||||
temp = _obj_for(calls, "observation.temp")
|
||||
assert type(temp).__name__ == "DummyScalar"
|
||||
assert temp.value == pytest.approx(10.0)
|
||||
assert float(temp.value) == pytest.approx(10.0)
|
||||
|
||||
img = _obj_for(calls, "observation.gray")
|
||||
assert type(img).__name__ == "DummyImage"
|
||||
assert type(img).__name__ == "DummyDepthImage" # single-channel -> DepthImage
|
||||
assert img.arr.shape == (8, 8, 1) # remains HWC
|
||||
assert _kwargs_for(calls, "observation.gray").get("static", False) is True
|
||||
|
||||
a = _obj_for(calls, "action.a")
|
||||
assert type(a).__name__ == "DummyScalar"
|
||||
assert a.value == pytest.approx(1.0)
|
||||
assert float(a.value) == pytest.approx(1.0)
|
||||
|
||||
# Blueprint sent once, with a spatial view for the image and time-series views for scalars
|
||||
assert len(blueprints) == 1
|
||||
bp = blueprints[0]
|
||||
assert {v.origin for v in _views_by_kind(bp, "Spatial2DView")} == {"observation.gray"}
|
||||
ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
|
||||
assert ts_views["observation"].contents == ["observation.temp"]
|
||||
assert ts_views["action"].contents == ["action.a"]
|
||||
|
||||
|
||||
def test_log_rerun_data_blueprint_sent_only_once(mock_rerun):
|
||||
"""The blueprint is built from the first call and not resent on subsequent calls."""
|
||||
vu, calls, blueprints = mock_rerun
|
||||
|
||||
vu.log_rerun_data(observation={"temp": 1.0}, action={"a": 2.0})
|
||||
assert len(blueprints) == 1
|
||||
first_blueprint = vu.log_rerun_data.blueprint
|
||||
|
||||
vu.log_rerun_data(observation={"temp": 3.0}, action={"a": 4.0})
|
||||
# Still only one blueprint, and the cached one is unchanged.
|
||||
assert len(blueprints) == 1
|
||||
assert vu.log_rerun_data.blueprint is first_blueprint
|
||||
|
||||
Reference in New Issue
Block a user