hangX/lerobot
1
0
Fork 0
forked from tangger/lerobot
Code Pull Requests Activity

format

Browse Source
This commit is contained in:
Quentin Gallouédec
2024-04-25 11:46:17 +02:00
parent fda092d5bf
commit 1ffc0e0d94
1 changed files with 10 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
lerobot/scripts/train.py
View File

@@ -280,21 +280,13 @@ def train(cfg: dict, out_dir=None, job_name=None):
# Temporary hack to move optimizer out of policy # Temporary hack to move optimizer out of policy
if isinstance(policy, ActPolicy): if isinstance(policy, ActPolicy):
optimizer_params_dicts = [ optimizer_params_dicts = [
{"params": [p for n, p in policy.named_parameters() if not n.startswith("backbone") and p.requires_grad]},
{ {
"params": [ "params": [p for n, p in policy.named_parameters() if n.startswith("backbone") and p.requires_grad],
p for n, p in policy.named_parameters() if not n.startswith("backbone") and p.requires_grad
]
},
{
"params": [
p for n, p in policy.named_parameters() if n.startswith("backbone") and p.requires_grad
],
"lr": policy.cfg.lr_backbone, "lr": policy.cfg.lr_backbone,
}, },
] ]
optimizer = torch.optim.AdamW( optimizer = torch.optim.AdamW(optimizer_params_dicts, lr=policy.cfg.lr, weight_decay=policy.cfg.weight_decay)
optimizer_params_dicts, lr=policy.cfg.lr, weight_decay=policy.cfg.weight_decay
)
lr_scheduler = None lr_scheduler = None
elif isinstance(policy, DiffusionPolicy): elif isinstance(policy, DiffusionPolicy):
optimizer = torch.optim.Adam( optimizer = torch.optim.Adam(
@@ -313,8 +305,6 @@ def train(cfg: dict, out_dir=None, job_name=None):
last_epoch=global_step - 1, last_epoch=global_step - 1,
) )
num_learnable_params = sum(p.numel() for p in policy.parameters() if p.requires_grad) 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()) num_total_params = sum(p.numel() for p in policy.parameters())
@@ -397,9 +387,7 @@ def train(cfg: dict, out_dir=None, job_name=None):
# create dataloader for online training # create dataloader for online training
concat_dataset = torch.utils.data.ConcatDataset([offline_dataset, online_dataset]) concat_dataset = torch.utils.data.ConcatDataset([offline_dataset, online_dataset])
weights = [1.0] * len(concat_dataset) weights = [1.0] * len(concat_dataset)
sampler = torch.utils.data.WeightedRandomSampler( sampler = torch.utils.data.WeightedRandomSampler(weights, num_samples=len(concat_dataset), replacement=True)
weights, num_samples=len(concat_dataset), replacement=True
)
dataloader = torch.utils.data.DataLoader( dataloader = torch.utils.data.DataLoader(
concat_dataset, concat_dataset,
num_workers=4, num_workers=4,