Add common, refactor eval with eval_policy

lerobot/scripts/train.py

@@ -2,7 +2,10 @@ import hydra
 import torch
 from termcolor import colored
 
-from ..lib.utils import set_seed
+from lerobot.common.envs.factory import make_env
+from lerobot.common.tdmpc import TDMPC
+
+from ..common.utils import set_seed
 
 
 @hydra.main(version_base=None, config_name="default", config_path="../configs")
@@ -11,6 +14,24 @@ def train(cfg: dict):
     set_seed(cfg.seed)
     print(colored("Work dir:", "yellow", attrs=["bold"]), cfg.log_dir)
 
+    env = make_env(cfg)
+    agent = TDMPC(cfg)
+    # ckpt_path = "/home/rcadene/code/fowm/logs/xarm_lift/all/default/2/models/offline.pt"
+    ckpt_path = "/home/rcadene/code/fowm/logs/xarm_lift/all/default/2/models/final.pt"
+    agent.load(ckpt_path)
+
+    # online training
+
+    eval_metrics = train_agent(
+        env,
+        agent,
+        num_episodes=10,
+        save_video=True,
+        video_dir=Path("tmp/2023_01_29_xarm_lift_final"),
+    )
+
+    print(eval_metrics)
+
 
 if __name__ == "__main__":
     train()