Merge remote-tracking branch 'upstream/main' into refactor_dp

lerobot/common/policies/factory.py

@@ -1,3 +1,10 @@
+import inspect
+
+from omegaconf import OmegaConf
+
+from lerobot.common.utils import get_safe_torch_device
+
+
 def make_policy(cfg):
     if cfg.policy.name == "tdmpc":
         from lerobot.common.policies.tdmpc.policy import TDMPCPolicy
@@ -19,10 +26,22 @@ def make_policy(cfg):
             **cfg.policy,
         )
     elif cfg.policy.name == "act":
-        from lerobot.common.policies.act.policy import ActionChunkingTransformerPolicy
+        from lerobot.common.policies.act.configuration_act import ActionChunkingTransformerConfig
+        from lerobot.common.policies.act.modeling_act import ActionChunkingTransformerPolicy
 
-        policy = ActionChunkingTransformerPolicy(cfg.policy, cfg.device)
-        policy.to(cfg.device)
+        expected_kwargs = set(inspect.signature(ActionChunkingTransformerConfig).parameters)
+        assert set(cfg.policy).issuperset(
+            expected_kwargs
+        ), f"Hydra config is missing arguments: {set(cfg.policy).difference(expected_kwargs)}"
+        policy_cfg = ActionChunkingTransformerConfig(
+            **{
+                k: v
+                for k, v in OmegaConf.to_container(cfg.policy, resolve=True).items()
+                if k in expected_kwargs
+            }
+        )
+        policy = ActionChunkingTransformerPolicy(policy_cfg)
+        policy.to(get_safe_torch_device(cfg.device))
     else:
         raise ValueError(cfg.policy.name)