mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-28 12:15:59 +00:00
Khalil Meftah
80 lines
2.8 KiB
Python
80 lines
2.8 KiB
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 sys
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock
|
|
|
|
import torch
|
|
|
|
import lerobot.policies.factory as policy_factory
|
|
|
|
|
|
def test_make_policy_keeps_peft_adapter_and_base_revisions_separate(monkeypatch):
|
|
cfg = SimpleNamespace(
|
|
type="mock",
|
|
device="cpu",
|
|
pretrained_path="user/adapter",
|
|
pretrained_revision="adapter-sha",
|
|
use_peft=True,
|
|
input_features={},
|
|
output_features={},
|
|
)
|
|
dataset_meta = SimpleNamespace(features={}, stats={})
|
|
|
|
base_policy = torch.nn.Linear(1, 1)
|
|
policy_from_pretrained = MagicMock(return_value=base_policy)
|
|
policy_class = SimpleNamespace(from_pretrained=policy_from_pretrained)
|
|
monkeypatch.setattr(policy_factory, "get_policy_class", lambda _: policy_class)
|
|
monkeypatch.setattr(policy_factory, "dataset_to_policy_features", lambda _: {})
|
|
monkeypatch.setattr(policy_factory, "validate_visual_features_consistency", lambda *args: None)
|
|
|
|
peft_config = SimpleNamespace(
|
|
base_model_name_or_path="user/base-policy",
|
|
revision="base-sha",
|
|
)
|
|
peft_config_from_pretrained = MagicMock(return_value=peft_config)
|
|
adapted_policy = torch.nn.Linear(1, 1)
|
|
peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"peft",
|
|
SimpleNamespace(
|
|
PeftConfig=SimpleNamespace(from_pretrained=peft_config_from_pretrained),
|
|
PeftModel=SimpleNamespace(from_pretrained=peft_model_from_pretrained),
|
|
),
|
|
)
|
|
|
|
policy = policy_factory.make_policy(cfg, ds_meta=dataset_meta)
|
|
|
|
assert policy is adapted_policy
|
|
peft_config_from_pretrained.assert_called_once_with(
|
|
"user/adapter",
|
|
revision="adapter-sha",
|
|
)
|
|
policy_from_pretrained.assert_called_once_with(
|
|
config=cfg,
|
|
dataset_stats=dataset_meta.stats,
|
|
dataset_meta=dataset_meta,
|
|
pretrained_name_or_path="user/base-policy",
|
|
revision="base-sha",
|
|
)
|
|
peft_model_from_pretrained.assert_called_once_with(
|
|
base_policy,
|
|
"user/adapter",
|
|
config=peft_config,
|
|
revision="adapter-sha",
|
|
is_trainable=True,
|
|
)
|