解决无法在Lerobot-Seed分支上多GPU训练问题

由于我购买的机械臂不是So-Arm101系列,而是由SeedStudio公司开发的STarAI机械臂,这款机械臂也是基于lerobot开发的。目前Lerobot主仓库是可以支持多GPU训练的,在官网https://huggingface.co/docs/lerobot/index中可以查看:

image

我的训练指令如下:

accelerate launch --num_processes=2 $(which lerobot-train) \
  --dataset.repo_id=/mnt/sdb/dataset/zzh/ACT/record-test-new1 \
  --policy.type=act \
  --policy.device=cuda \
  --policy.chunk_size=100 \
  --policy.n_action_steps=50 \
  --policy.use_amp=true \
  --policy.repo_id=starai/my_policy \
  --batch_size=32 \
  --policy.optimizer_lr=2e-05 \
  --num_workers=4 \
  --output_dir=outputs/train/act_viola_test1 \
  --job_name=act_viola_test \
  --wandb.enable=False \
  --steps=100000 \
  --save_checkpoint=True \
  --save_freq=50000

但是存在报错,排查了一下原因发现报错和多 GPU 本身无关,而是配置和 ACTPolicy 构造函数不匹配(不太清除跟其他policy有没有产生冲突,我还没做到那一步)
原因是在lerobot-main/src/lerobot/policies/factory.py中的476行传递了2个参数给策略类
image

而我们翻看lerobot-main/src/lerobot/policies/act/modeling_act.py中的__init__函数时:
image

发现ACTPolicy.init()确实只接受一个 config参数,不接受这两个额外参数。
我推测原因是因为Seed分支落后与Lerobot主分支,相关的代码还没来得及适配,我询问了技术支持,证实了这个想法

解决办法:
很简单,既然factory传入了两个额外的参数,那我只需要在ACTPolicy.init()函数后加入2个形参,然后不去管它就好了
将函数头修改为:

    def __init__(
        self,
        config: ACTConfig,
        dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
        dataset_meta: Any | None = None,
    ):

并且导入库

from typing import Any

再次启动训练指令,嗯!错误消失,可以愉快的训练了。

image

posted @ 2026-03-09 12:26  zzzking778  阅读(8)  评论(0)    收藏  举报