tangger/lerobot
1
0
Fork 1
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

added optimizer and sac to factory.py

Browse Source
This commit is contained in:
Michel Aractingi
2024-12-23 14:12:03 +01:00
parent 1e2a757cd3
commit b0e2fcdba7
1 changed files with 9 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
lerobot/scripts/train.py
View File

@@ -93,6 +93,15 @@ def make_optimizer_and_scheduler(cfg, policy):
elif policy.name == "tdmpc": elif policy.name == "tdmpc":
optimizer = torch.optim.Adam(policy.parameters(), cfg.training.lr) optimizer = torch.optim.Adam(policy.parameters(), cfg.training.lr)
lr_scheduler = None lr_scheduler = None
elif policy.name == "sac":
optimizer = torch.optim.Adam([
{'params': policy.actor.parameters(), 'lr': policy.config.actor_lr},
{'params': policy.critic_ensemble.parameters(), 'lr': policy.config.critic_lr},
{'params': policy.temperature.parameters(), 'lr': policy.config.temperature_lr},
])
lr_scheduler = None
elif cfg.policy.name == "vqbet": elif cfg.policy.name == "vqbet":
from lerobot.common.policies.vqbet.modeling_vqbet import VQBeTOptimizer, VQBeTScheduler from lerobot.common.policies.vqbet.modeling_vqbet import VQBeTOptimizer, VQBeTScheduler