wangzixiang
/
lerobot
mirror of https://github.com/huggingface/lerobot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Make sure to make remove all traces of omegaconf from policy config

This commit is contained in:
Alexander Soare 2024-04-15 09:59:18 +01:00
parent 9241b5e830
commit 40d417ef60
1 changed files with 9 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
lerobot/common/policies/factory.py
View File

@ -1,5 +1,7 @@
import inspect import inspect
from omegaconf import OmegaConf
from lerobot.common.utils import get_safe_torch_device from lerobot.common.utils import get_safe_torch_device
@ -33,7 +35,13 @@ def make_policy(cfg):
assert set(cfg.policy).issuperset( assert set(cfg.policy).issuperset(
expected_kwargs expected_kwargs
), f"Hydra config is missing arguments: {set(cfg.policy).difference(expected_kwargs)}" ), f"Hydra config is missing arguments: {set(cfg.policy).difference(expected_kwargs)}"
policy_cfg = ActConfig(**{k: v for k, v in cfg.policy.items() if k in expected_kwargs}) policy_cfg = ActConfig(
**{
k: v
for k, v in OmegaConf.to_container(cfg.policy, resolve=True).items()
if k in expected_kwargs
}
)
policy = ActPolicy(policy_cfg) policy = ActPolicy(policy_cfg)
policy.to(get_safe_torch_device(cfg.device)) policy.to(get_safe_torch_device(cfg.device))
else: else: