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

README.md

@@ -89,9 +89,6 @@ eval_episodes=7
 
 **style**
 ```
-isort lerobot
+isort lerobot && isort test && black lerobot && black test
-black lerobot
+pylint lerobot && pylint test  # not enforce for now
-isort test
-black test
-pylint lerobot
 ```

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)
+        policy = make_policy(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 = 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)
+    policy = make_policy(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)
 
     td_policy = TensorDictModule(
         policy,

test/test_envs.py

@@ -6,6 +6,8 @@ from lerobot.common.envs.factory import make_env
 from lerobot.common.envs.pusht import PushtEnv
 from lerobot.common.envs.simxarm import SimxarmEnv
 
+from .utils import init_config
+
 
 def print_spec_rollout(env):
     print("observation_spec:", env.observation_spec)
@@ -83,14 +85,6 @@ def test_pusht(from_pixels, pixels_only):
     ],
 )
 def test_factory(config_name):
-    import hydra
+    cfg = init_config(config_name)
-    from hydra import compose, initialize
-
-    config_path = "../lerobot/configs"
-    hydra.core.global_hydra.GlobalHydra.instance().clear()
-    initialize(config_path=config_path)
-    cfg = compose(config_name=config_name)
-
     env = make_env(cfg)
-
     check_env_specs(env)

test/test_policies.py

@@ -0,0 +1,17 @@
+import pytest
+
+from lerobot.common.policies.factory import make_policy
+
+from .utils import init_config
+
+
+@pytest.mark.parametrize(
+    "config_name",
+    [
+        "default",
+        "pusht",
+    ],
+)
+def test_factory(config_name):
+    cfg = init_config(config_name)
+    policy = make_policy(cfg)

test/utils.py

@@ -0,0 +1,11 @@
+import hydra
+from hydra import compose, initialize
+
+CONFIG_PATH = "../lerobot/configs"
+
+
+def init_config(config_name):
+    hydra.core.global_hydra.GlobalHydra.instance().clear()
+    initialize(config_path=CONFIG_PATH)
+    cfg = compose(config_name=config_name)
+    return cfg