Add policies/factory, Add test, Add _self_ in config

lerobot/common/policies/factory.py

@@ -0,0 +1,21 @@
+from lerobot.common.policies.tdmpc import TDMPC
+
+
+def make_policy(cfg):
+    if cfg.policy == "tdmpc":
+        policy = TDMPC(cfg)
+    else:
+        raise ValueError(cfg.policy)
+
+    if cfg.pretrained_model_path:
+        # TODO(rcadene): hack for old pretrained models from fowm
+        if cfg.policy == "tdmpc" and "fowm" in cfg.pretrained_model_path:
+            if "offline" in cfg.pretrained_model_path:
+                policy.step[0] = 25000
+            elif "final" in cfg.pretrained_model_path:
+                policy.step[0] = 100000
+            else:
+                raise NotImplementedError()
+        policy.load(cfg.pretrained_model_path)
+
+    return policy

lerobot/common/policies/tdmpc.py

@@ -5,7 +5,7 @@ import numpy as np
 import torch
 import torch.nn as nn
 
-import lerobot.common.tdmpc_helper as h
+import lerobot.common.policies.tdmpc_helper as h
 
 
 class TOLD(nn.Module):

lerobot/configs/default.yaml

@@ -40,6 +40,7 @@ state_dim: 4
 action_dim: 4
 
 # TDMPC
+policy: tdmpc
 
 # planning
 mpc: true

lerobot/configs/pusht.yaml

@@ -1,4 +1,5 @@
 defaults:
+  - _self_
   - default
 
 hydra:

lerobot/scripts/eval.py

@@ -10,7 +10,7 @@ from termcolor import colored
 from torchrl.envs import EnvBase
 
 from lerobot.common.envs.factory import make_env
-from lerobot.common.tdmpc import TDMPC
+from lerobot.common.policies.factory import make_policy
 from lerobot.common.utils import set_seed
 
 
@@ -111,15 +111,7 @@ def eval(cfg: dict, out_dir=None):
     env = make_env(cfg)
 
     if cfg.pretrained_model_path:
-        policy = TDMPC(cfg)
-        if "offline" in cfg.pretrained_model_path:
-            policy.step[0] = 25000
-        elif "final" in cfg.pretrained_model_path:
-            policy.step[0] = 100000
-        else:
-            raise NotImplementedError()
-        policy.load(cfg.pretrained_model_path)
-
+        policy = make_policy(cfg)
         policy = TensorDictModule(
             policy,
             in_keys=["observation", "step_count"],

lerobot/scripts/train.py

@@ -11,10 +11,9 @@ from torchrl.data.datasets.openx import OpenXExperienceReplay
 from torchrl.data.replay_buffers import PrioritizedSliceSampler
 
 from lerobot.common.datasets.factory import make_offline_buffer
-from lerobot.common.datasets.simxarm import SimxarmExperienceReplay
 from lerobot.common.envs.factory import make_env
 from lerobot.common.logger import Logger
-from lerobot.common.tdmpc import TDMPC
+from lerobot.common.policies.factory import make_policy
 from lerobot.common.utils import set_seed
 from lerobot.scripts.eval import eval_policy
 
@@ -51,17 +50,7 @@ def train(cfg: dict, out_dir=None, job_name=None):
     print(colored("Work dir:", "yellow", attrs=["bold"]), out_dir)
 
     env = make_env(cfg)
-    policy = TDMPC(cfg)
-    if cfg.pretrained_model_path:
-        # TODO(rcadene): hack for old pretrained models from fowm
-        if "fowm" in cfg.pretrained_model_path:
-            if "offline" in cfg.pretrained_model_path:
-                policy.step[0] = 25000
-            elif "final" in cfg.pretrained_model_path:
-                policy.step[0] = 100000
-            else:
-                raise NotImplementedError()
-        policy.load(cfg.pretrained_model_path)
+    policy = make_policy(cfg)
 
     td_policy = TensorDictModule(
         policy,