File size: 1,321 Bytes
52f973d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 | """Load this repository offline after snapshot_download, using official LeRobot."""
from pathlib import Path
def load_policy(checkpoint_dir, device="cpu"):
from lerobot.configs.policies import PreTrainedConfig
from lerobot.policies.factory import get_policy_class, make_pre_post_processors
root = Path(checkpoint_dir).resolve()
config = PreTrainedConfig.from_pretrained(root)
config.device = device
config.pretrained_path = root
overrides = {"device_processor": {"device": device}}
if config.type == "smolvla":
assets = root / "vlm_assets"
config.vlm_model_name = str(assets)
config.load_vlm_weights = False
overrides["tokenizer_processor"] = {"tokenizer_name": str(assets)}
elif config.type == "act":
# The full trained backbone is in model.safetensors; avoid fetching its
# initialization weights before restoring that complete state dictionary.
config.pretrained_backbone_weights = None
policy = get_policy_class(config.type).from_pretrained(root, config=config).eval().to(device)
pre, post = make_pre_post_processors(
config,
pretrained_path=root,
preprocessor_overrides=overrides,
postprocessor_overrides={"device_processor": {"device": "cpu"}},
)
return policy, pre, post
|