rename make_optimizer

Author: aliberts

lerobot/scripts/train.py

@@ -25,7 +25,7 @@
 from lerobot.scripts.eval import eval_policy
 
 
-def make_optimizer(cfg, policy):
+def make_optimizer_and_scheduler(cfg, policy):
     if cfg.policy.name == "act":
         optimizer_params_dicts = [
             {
@@ -321,7 +321,7 @@ def train(cfg: dict, out_dir=None, job_name=None):
 
     # Create optimizer and scheduler
     # Temporary hack to move optimizer out of policy
-    optimizer, lr_scheduler = make_optimizer(cfg, policy)
+    optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
 
     num_learnable_params = sum(p.numel() for p in policy.parameters() if p.requires_grad)
     num_total_params = sum(p.numel() for p in policy.parameters())

tests/scripts/save_policy_to_safetensor.py

@@ -7,7 +7,7 @@
 from lerobot.common.datasets.factory import make_dataset
 from lerobot.common.policies.factory import make_policy
 from lerobot.common.utils.utils import init_hydra_config, set_global_seed
-from lerobot.scripts.train import make_optimizer
+from lerobot.scripts.train import make_optimizer_and_scheduler
 from tests.utils import DEFAULT_CONFIG_PATH
 
 
@@ -25,7 +25,7 @@ def get_policy_stats(env_name, policy_name, extra_overrides=None):
     dataset = make_dataset(cfg)
     policy = make_policy(cfg, dataset_stats=dataset.stats)
     policy.train()
-    optimizer, _ = make_optimizer(cfg, policy)
+    optimizer, _ = make_optimizer_and_scheduler(cfg, policy)
 
     dataloader = torch.utils.data.DataLoader(
         dataset,