Refactor configs to have env in seperate yaml + Fix training

2024-02-25 17:42:47 +00:00
parent eec134d72b
commit b16c334825
13 changed files with 146 additions and 54 deletions
--- a/lerobot/common/policies/factory.py
+++ b/lerobot/common/policies/factory.py
@@ -1,9 +1,12 @@
-from lerobot.common.policies.tdmpc import TDMPC
-
-
 def make_policy(cfg):
    if cfg.policy.name == "tdmpc":
+        from lerobot.common.policies.tdmpc import TDMPC
+
        policy = TDMPC(cfg.policy)
+    elif cfg.policy.name == "diffusion":
+        from lerobot.common.policies.diffusion import DiffusionPolicy
+
+        policy = DiffusionPolicy(cfg.policy)
    else:
        raise ValueError(cfg.policy.name)