Clean up loading code

- Centralized instantiation of the PEFT wrapper in `make_policy` for inference
  (e.g. in `lerobot-record`)
- Training a PEFT policy also sets `cfg.use_peft` so that all inference code loading
  the policy can rely on that attribute to identify if PEFT loading is needed
- Modified RTC example to also include PEFT policies. Mostly because this is an example
  I'm currently exploring.
This commit is contained in:
nemo
2025-11-24 15:30:26 +01:00
parent 841c76c7d3
commit e9b3889bd2
4 changed files with 28 additions and 22 deletions
+8 -1
View File
@@ -455,7 +455,14 @@ def demo_cli(cfg: RTCDemoConfig):
if cfg.policy.type == "pi05" or cfg.policy.type == "pi0":
config.compile_model = cfg.use_torch_compile
policy = policy_class.from_pretrained(cfg.policy.pretrained_path, config=config)
if config.use_peft:
from peft import PeftModel
policy = policy_class(config=config)
policy = PeftModel.from_pretrained(policy, cfg.policy.pretrained_path)
else:
policy = policy_class.from_pretrained(cfg.policy.pretrained_path, config=config)
# Turn on RTC
policy.config.rtc_config = cfg.rtc