| """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": |
| |
| |
| 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 |
|
|
