mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 18:26:11 +00:00
pepijn
9d00293f49
Package fitted tokenizer artifacts with processor pipelines so pretrained PI052 checkpoints restore their saved recipe and normalization state without refitting. Co-authored-by: Cursor <cursoragent@cursor.com>
123 lines
4.0 KiB
Python
123 lines
4.0 KiB
Python
#!/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.
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from lerobot.policies.pi052.fit_fast_tokenizer import (
|
|
_apply_relative_actions,
|
|
_dataset_signature,
|
|
_is_global_leader,
|
|
_normalize_actions,
|
|
_select_episode_indices,
|
|
_validate_fast_reconstruction,
|
|
)
|
|
|
|
|
|
def test_fast_tokenizer_fit_uses_training_mean_std_normalization():
|
|
actions = np.array([[[1.0, 7.0], [3.0, 3.0]]], dtype=np.float32)
|
|
stats = {"mean": [2.0, 5.0], "std": [0.5, 2.0]}
|
|
|
|
normalized = _normalize_actions(actions, "MEAN_STD", stats)
|
|
|
|
np.testing.assert_allclose(normalized, [[[-2.0, 1.0], [2.0, -1.0]]])
|
|
|
|
|
|
def test_fast_tokenizer_fit_quantiles_match_training_without_clipping():
|
|
actions = np.array([[[-1.0], [3.0]]], dtype=np.float32)
|
|
stats = {"q01": [0.0], "q99": [2.0]}
|
|
|
|
normalized = _normalize_actions(actions, "QUANTILES", stats)
|
|
|
|
np.testing.assert_allclose(normalized, [[[-2.0], [2.0]]])
|
|
|
|
|
|
def test_fast_tokenizer_cache_signature_tracks_stats_and_episode_selection():
|
|
kwargs = {
|
|
"dataset_repo_id": "org/dataset",
|
|
"base_tokenizer_name": "physical-intelligence/fast",
|
|
"n_samples": 100,
|
|
"chunk_size": 20,
|
|
"normalization_mode": "QUANTILES",
|
|
"dataset_revision": "main",
|
|
"episodes": [1, 2, 3],
|
|
"exclude_episodes": [2],
|
|
"use_relative_actions": False,
|
|
"relative_action_mask": None,
|
|
}
|
|
|
|
first = _dataset_signature(**kwargs, action_stats={"q01": [0.0], "q99": [1.0]})
|
|
changed_stats = _dataset_signature(**kwargs, action_stats={"q01": [0.0], "q99": [2.0]})
|
|
changed_selection = _dataset_signature(
|
|
**{**kwargs, "exclude_episodes": [2, 3]},
|
|
action_stats={"q01": [0.0], "q99": [1.0]},
|
|
)
|
|
|
|
assert first != changed_stats
|
|
assert first != changed_selection
|
|
|
|
|
|
def test_fast_tokenizer_uses_only_global_rank_zero(monkeypatch):
|
|
monkeypatch.setenv("RANK", "8")
|
|
monkeypatch.setenv("LOCAL_RANK", "0")
|
|
assert not _is_global_leader()
|
|
|
|
monkeypatch.setenv("RANK", "0")
|
|
assert _is_global_leader()
|
|
|
|
|
|
def test_fast_tokenizer_episode_selection_applies_allowlist_and_exclusions():
|
|
selected = _select_episode_indices([0, 1, 2, 3], episodes=[1, 2, 3], exclude_episodes=[2])
|
|
|
|
assert selected == [1, 3]
|
|
|
|
|
|
def test_fast_tokenizer_relative_actions_match_training_transform():
|
|
actions = np.array([[[2.0, 10.0], [3.0, 11.0]]], dtype=np.float32)
|
|
states = np.array([[1.0, 4.0]], dtype=np.float32)
|
|
|
|
relative = _apply_relative_actions(actions, states, [True, False])
|
|
|
|
np.testing.assert_allclose(relative, [[[1.0, 10.0], [2.0, 11.0]]])
|
|
|
|
|
|
class _RoundTripTokenizer:
|
|
def __init__(self, offset: float = 0.0):
|
|
self.offset = offset
|
|
|
|
def __call__(self, actions):
|
|
return actions
|
|
|
|
def decode(self, tokens):
|
|
return tokens + self.offset
|
|
|
|
|
|
def test_fast_tokenizer_reconstruction_validation_reports_error():
|
|
actions = np.arange(24, dtype=np.float32).reshape(2, 3, 4) / 24
|
|
|
|
report, decoded = _validate_fast_reconstruction(_RoundTripTokenizer(0.05), actions, 0.1, 0.1)
|
|
|
|
np.testing.assert_allclose(decoded, actions + 0.05)
|
|
assert report["reconstruction_rmse"] == pytest.approx(0.05)
|
|
assert report["max_dim_rmse"] == pytest.approx(0.05)
|
|
|
|
|
|
def test_fast_tokenizer_reconstruction_validation_rejects_large_error():
|
|
actions = np.arange(24, dtype=np.float32).reshape(2, 3, 4) / 24
|
|
|
|
with pytest.raises(RuntimeError, match="exceeds the configured limit"):
|
|
_validate_fast_reconstruction(_RoundTripTokenizer(0.25), actions, 0.1, 0.2)
|