Add VQ-BeT (#166)

2024-06-26 16:55:02 +09:00
parent 9aa4cdb976
commit 7d1542cae1
9 changed files with 2644 additions and 1 deletions
--- a/lerobot/scripts/train.py
+++ b/lerobot/scripts/train.py
@@ -88,6 +88,11 @@ def make_optimizer_and_scheduler(cfg, policy):
    elif policy.name == "tdmpc":
        optimizer = torch.optim.Adam(policy.parameters(), cfg.training.lr)
        lr_scheduler = None
+    elif cfg.policy.name == "vqbet":
+        from lerobot.common.policies.vqbet.modeling_vqbet import VQBeTOptimizer, VQBeTScheduler
+
+        optimizer = VQBeTOptimizer(policy, cfg)
+        lr_scheduler = VQBeTScheduler(optimizer, cfg)
    else:
        raise NotImplementedError()