robomimic 的 Lift-PH 复现系列终于写到阶段 7这是整个项目里视觉模仿学习最关键的一段Vision BC视觉行为克隆从数据构建、视觉编码、正式训练到故障恢复和最终评估的完整链路。前几个阶段如果已经把环境、数据集和低维 BC 跑通那么这一阶段就是把策略从“只看状态向量”升级到“真正看懂图像”难度跨了一个台阶但效果也最贴近真实机器人操作。这一篇我按照自己实际跑通 Lift-PH Vision BC 的过程整理重点不是贴文档而是回答那些文档里不写、但你跑起来一定会遇到的问题图像观测怎么从 HDF5 里取出来、ResNet 视觉编码器怎么和动作头拼在一起、训练中断了怎么优雅续上、成功率到底怎么算。适合正在复现 robomimic、想把视觉模仿学习真正跑起来的同学也适合已经跑通低维 BC、准备往视觉方向进阶的读者。1. 阶段全貌Vision BC 在 Lift-PH 中的定位与复现链路1.1 Lift-PH 是什么为什么阶段 7 要交给 Vision BCLift 是 robomimic 里最经典的仿真任务一个机械臂需要从桌面上抓起一个方块并抬升到目标高度。任务本身感知上不复杂但它是验证模仿学习算法的“最小可用环境”——动作空间小、状态变化直观、失败模式容易排查。后缀 PH 指的是 Proficient Human表示这批演示数据来自熟练人类操作者的遥操作记录而不是脚本或低质量随机策略。PH 数据的质量是后续一切训练的前提robomimic 官方把 Lift-PH 作为一个标准基准几乎所有算法BC、BC-RNN、IRL、CQL 等都在这个任务上有公开对比。阶段 7 之所以要引入 Vision BC是因为真实机器人场景里你不可能总是拿到完美的低维状态比如夹爪位置、物体坐标更多时候只有一个摄像头画面。Vision BC 直接从图像观测学习动作映射是打通“感知-控制”最简单直接的一步。它不需要强化学习那样的在线交互只需要离线演示数据训练完成后用 rollout 去验证策略好坏非常适合作为视觉模仿学习的第一个实践项目。1.2 从数据到部署的一条完整链路阶段 7 并不只是运行一行训练命令那么简单前面有数据构建后面有故障恢复和评估整条链路是确认演示数据Lift-PH 的 HDF5 文件里有哪些字段、图像是什么尺寸、动作空间是几维。构建训练数据集把 HDF5 中的原始观测转换成 robomimic Dataset 所需的观测字典配置 obs_keys、图像归一化、图像增强策略。设计视觉编码器搭建 ResNet 骨干、空间 softmax、与低维状态/动作的融合方式。正式训练写好 config 文件启动 train.py监控 action loss 和验证集表现。故障恢复训练中断后利用 checkpoint 续训处理显存不足和数据异常。评估与验证用 run_trained_agent.py 做 rollout计算成功率分析失败模式。这一步和前面低维 BC 最大的区别在于图像观测的引入让数据 pipeline 和网络结构都复杂了。图像数据体积大读取/预处理容易成为瓶颈视觉编码器参数量大显存和训练时间的压力也明显上升。所以阶段 7 的“数据构建”和“视觉编码”两个环节其实是整个 Vision BC 能否顺利训练的关键值得专门拆开讲。2. 数据构建把 HDF5 演示变成模型能吃的“视觉动作”样本2.1 先看清数据集结构别跳过这一步robomimic 的演示数据统一存放在 HDF5 文件里。Lift-PH 的官方路径一般是datasets/lift_ph_image.hdf5但我建议你用自己的数据目录避免和官方默认路径搞混。拿到数据集后第一件事不是急着训练而是用脚本检查它的内部结构import h5py dataset_path ../datasets/lift_ph_image.hdf5 with h5py.File(dataset_path, r) as f: print(root keys:, list(f.keys())) demo_keys list(f[data].keys()) print(num demos:, len(demo_keys)) demo f[data/demo_0] print(obs keys:, list(demo[obs].keys())) print(actions shape:, demo[actions].shape) print(image shape:, demo[obs][agentview_image].shape)这个检查非常关键。我遇到过不止一次数据集版本不对、obs key 命名不同导致后面训练时报 KeyError。Lift-PH 的agentview_image和eye_in_hand_image是常用的两个视觉观测shape 一般是(T, 3, H, W)通道在前或(T, H, W, 3)robomimic 会按 config 里设定的方式处理。动作维度 Lift 一般是 7 维即末端位置 3 维 四元数姿态 4 维有的数据版本会额外包含夹爪指令那就是 8 维或 10 维。你要在 check 时确认actions的维度然后在 config 里把ac_dim配成对应值。2.2 用 dataset_states_to_obs_dict_replay 完成观测转换robomimic 官方训练脚本里数据加载通常不直接操作 HDF5而是通过SequenceDataset配合一个 obs 字典转换函数完成。在 Vision BC 的配置里你需要明确告诉 robomimic哪些 key 来自 demo 的obs组哪些 key 是图像、哪些是低维状态。比较省事的做法是在 config 里配置data.obs_keys然后让 robomimic 内部完成转换。不过如果你的数据不是官方格式或者你想增加自定义特征就需要自己写一个转换函数def dataset_states_to_obs_dict_replay(demo): 将 hdf5 demo 中的原始 obs 转为 robomimic 训练所需的 obs 字典。 obs_dict { agentview_image: demo[obs][agentview_image], eye_in_hand_image: demo[obs][eye_in_hand_image], robot0_eef_pos: demo[obs][robot0_eef_pos], robot0_eef_quat: demo[obs][robot0_eef_quat], } return obs_dict然后传给 SequenceDataset 时指定这个函数。这样做的好处是你能完全控制哪些观测被送进网络。我在实际操作中习惯保留robot0_eef_pos和robot0_eef_quat这两个低维状态跟图像一起输入因为机械臂末端位姿对 Lift 这种需要精确到达物体的任务特别有用只靠图像会让策略学得更慢。2.3 图像像素预处理与归一化直接影响收敛速度很多人忽略图像归一化直接把 0-255 的 uint8 像素喂给 ResNet。这在 robomimic 里通常也能跑但收敛明显更慢而且 loss 容易出现奇怪的跳动。robomimic 提供data.image_obs_normalization之类的配置但更通用的做法是在 Dataset 的 obs 处理里对图像做 ImageNet 归一化transform torchvision.transforms.Normalize( mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] )注意归一化操作要在把 uint8 转成 float 并缩放到 [0,1] 之后进行。如果你用的是 robomimic 的ImageNormalize工具它会自动处理这些但自己写转换函数时一定要记得这一步。经历过一次效果对比后我的结论是归一化对 Vision BC 来说不是可选项而是必选项。它让 ResNet 的预训练权重如果用了发挥正常作用也避免不同光照条件下像素值偏移对训练造成干扰。2.4 数据质量自检清单数据是模仿学习的上限。Lift-PH 演示数据虽然叫“Proficient Human”但不代表每个 demo 都是完美可用的。我在训练前通常会做一套快速自检动作序列是否有长时间静止片段若有可能是遥操作故障这类 demo 会让策略学到“什么都不做”。图像是否出现全黑/全白帧这些帧会污染视觉编码器最好直接滤除。action 数值分布是否合理比如连续动作是否有 NaN末端位置是否突然跳变。所有 demo 的 obs key 是否一致不同 demo 缺少某个 key 是训练崩溃常见原因。建议写一个几行的检查脚本把每个 demo 的 actions 和图像帧做统计把异常的 demo id 打印出来。这一步虽然不产生“成果”但能在训练前拦截大量问题省掉后面排查的几小时。3. 视觉编码ResNet18、空间 softmax 与多视角融合3.1 为什么 Lift 场景必须看图像如果你已经跑过 Lift-PH 的低维 BC会发现只用robot0_eef_pos、robot0_eef_quat、物体位置这类状态也能学到不错的成功率。那 Vision BC 还有必要吗有而且价值体现在两点一是泛化低维状态在仿真里可以直接拿到但迁移到真实机器人时往往只能靠视觉二是信息密度图像里包含物体形状、位置、遮挡关系、夹爪开合状态这些隐含特征很难手工设计成低维向量。具体到 Lift 任务机械臂要先接近方块、调整夹爪方向、闭合夹爪、抬升。这个过程里“夹爪和物体的相对位置关系”是最核心的信息在图像里一目了然但在低维状态里需要精确的坐标变换才能表达。Vision BC 让网络自己从像素里提取这些关系省去了大量人工特征工程。3.2 视觉骨干网络搭建与参数选择robomimic 的 Vision BC 默认支持多种视觉骨干最常用的是 ResNet18。相比 ResNet50 或更重的网络ResNet18 在 Lift 这类简单视觉任务上完全够用训练更快显存占用更低对新手更友好。我的配置里一般这样设vision_encoder: { core_net_class: ResNet18Conv, core_net_kwargs: { input_channel: 3, feature_dim: 64, use_spatial_softmax: true, spatial_softmax_kwargs: { num_kp: 32, temperature: 1.0, learnable_temperature: true, output_size: 64 } } }这里use_spatial_softmax很关键。空间 softmax 会把卷积特征图上的每个位置转换成一个 2D 关键点坐标相当于让网络自己学会“关注图像里哪些位置”。对 Lift 来说这些关键点往往会落在方块边缘、夹爪指尖、目标区域附近可解释性很强。我建议在阶段 7 一定开启 spatial softmax它能让视觉特征更紧凑也更容易和后面的 MLP 层配合。3.3 视觉特征与低维状态特征的融合方式Vision BC 并不一定只用图像。前面我提到会保留robot0_eef_pos和robot0_eef_quat这时网络就需要处理两类输入视觉特征和低维状态。robomimic 的处理方式是把它们 concat 起来再送入后续的 MLP 层。具体来说图像经过 ResNet18 spatial softmax 输出一个feature_dim维的向量比如 64 维低维状态经过一个小的 MLP 编码成相同或不同维度的向量两者拼接后进入 actor 网络的隐层。这个设计的好处是视觉负责“看”状态负责“感知自身”两者互补。配置时注意actor_net.vision_encoder的输出维度要和低维状态编码后的维度对齐否则 concat 时会出问题。我在实际调试中的一个心得如果你发现网络训练正常但 rollout 时机械臂到处乱摸可以先试试只保留视觉输入去掉低维状态看成功率是否变化。这能帮你判断网络到底依赖哪路信息。如果去掉低维状态成功率骤降说明视觉编码还没学好如果成功率不变说明视觉已经掌握了核心信息低维状态只是辅助。3.4 图像增强的取舍别照搬分类任务的套路图像增强在视觉模仿学习里是个双刃剑。robomimic 支持随机裁剪、颜色抖动、翻转等增强方式。在 Lift 仿真环境里背景是固定的物体颜色也是固定的我实际测试下来随机裁剪有时反而会切掉机械臂末端导致策略困惑颜色抖动对收敛也没有明显帮助。我的建议是阶段 7 先关闭图像增强只做归一化。等 baseline 跑通之后再试着开一个轻微的随机裁剪比如裁剪后 resize 回原尺寸看看 val loss 和 rollout 成功率是否提升。不要在初期就引入复杂增强否则你会分不清是增强的问题还是网络结构的问题。仿真环境里固定的视觉分布本身就不是增强的用武之地增强真正发挥价值是在真实机器人数据跨场景迁移的时候。4. 正式训练config 逐项拆解与 loss 观测4.1 一份可用的 Vision BC 配置参考下面是我在 Lift-PH 上实际用过的精简配置去掉了与阶段 7 无关的冗余项保留核心参数{ experiment: { name: lift_ph_vbc, save: { enabled: true, frequency: 500 }, validate: { enabled: true, frequency: 10 } }, train: { num_epochs: 600, batch_size: 32, optimizer: { type: adamw, lr: 0.0003, weight_decay: 0.0001 }, output_dir: ../ckpt, checkpoint_frequency: 10, log_frequency: 100 }, algo: { share_encoder: false, actor_net: { ac_dim: 7, hidden_dim: 1024, layers: 3, vision_encoder: { core_net_class: ResNet18Conv, core_net_kwargs: { input_channel: 3, feature_dim: 64, use_spatial_softmax: true, spatial_softmax_kwargs: { num_kp: 32, temperature: 1.0, learnable_temperature: true, output_size: 64 } } } }, gmm: { enabled: true, num_modes: 5, min_std: 0.0001 } }, data: { dataset_path: ../datasets/lift_ph_image.hdf5, obs_keys: [agentview_image, eye_in_hand_image, robot0_eef_pos, robot0_eef_quat], image_obs_normalization: true, image_augmentation: false } }几个参数我说一下选择原因。batch_size32在单卡上处理双视角 ResNet18 是比较稳妥的值再大容易出现显存不足再小则训练不稳定。lr0.0003配 AdamW 是视觉主干网络常用的初始学习率如果 loss 震荡可以降到 0.0001。gmm.enabledtrue打开 GMM 输出头让网络预测动作的混合高斯分布这对模仿多模态演示比如同一种任务存在多种轨迹非常关键比直接回归动作效果好得多。4.2 训练启动与 loss 日志解读启动训练用 robomimic 的train.py把上面配置存成lift_ph_vision_bc.json然后执行python train.py \ --config configs/lift_ph_vision_bc.json \ --dataset ../datasets/lift_ph_image.hdf5训练日志会输出到output_dir下的lift_ph_vbc目录。重点关注两个指标train/action_loss和val/action_loss。前者是训练集上的动作预测误差后者是验证集上的误差。我观测到的一个典型过程是前 10 个 epochloss 从初始值快速下降val loss 也跟着降到 50-100 epoch 后train loss 继续缓慢下降val loss 开始波动甚至上升这时模型可能进入过拟合。robomimic 的 validate 设置每 10 个 epoch 评估一次你可以根据 val loss 的趋势决定是否提前终止而不是傻等到 600 epoch。如果显存和算力允许建议开 TensorBoard 或至少保存 csv 日志方便回看。4.3 显存占用与训练时长估算Vision BC 的显存压力比低维 BC 大得多尤其开了双视角图像。ResNet18 batch_size32 双视角在单张 12GB 显存卡上大概会占到 8-10GB属于比较紧的状态。如果显存不足优先做两件事减小 batch_size 到 16 或 8如果还不行就把eye_in_hand_image从 obs_keys 里去掉只留一个agentview_image。在 Lift 任务上手眼相机视角确实有帮助但少一个视角并不会让策略完全失效。训练时长方面600 epoch 在单张消费级显卡RTX 3080 级别上双视角 ResNet18 大约需要 1.5 到 3 小时。如果只是验证流程可以先设 50 epoch 跑通一遍确认日志、checkpoint、评估脚本都正常后再开长训练。我习惯的做法是先用 10 个 epoch “冒烟测试”检查代码和配置没有低级错误再正式跑长训练。5. 故障恢复中断续训、OOM 与脏数据急救手册5.1 用 checkpoint 回到断点别让训练从头再来训练跑了一半因为服务器重启、显存被占、或者手滑关掉终端导致中断是再常见不过的事。robomimic 默认会在output_dir下保存 checkpoint 文件命名类似checkpoint_epoch_50.pt、checkpoint_epoch_100.pt。恢复训练用--resume参数python train.py \ --config ../ckpt/lift_ph_vbc/config.json \ --resume ../ckpt/lift_ph_vbc/checkpoint_epoch_50.pt注意--resume会加载模型权重和优化器状态并继续从 checkpoint 记录的 epoch 开始训练。如果你只想加载模型权重但重新开始训练则不要用--resume而是手动设置train.start_epoch0并加载权重。这两个场景差别很大别搞混。我遇到过一次“续训后 loss 暴涨”的问题原因是续训时 config 里的 batch_size 被改了优化器状态和新的 batch 不匹配。恢复训练时最好保持和中断前完全一致的配置。5.2 显存不足的应急方案OOMOut of Memory是 Vision BC 训练中最常见的崩溃原因。如果你的日志里出现CUDA out of memory先别急着调网络结构按下面的顺序排查减小 batch_size从 32 到 16再到 8找到能跑稳的值。检查是否同时运行了多个占用显存的进程nvidia-smi看一下。如果板卡支持设置PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True再启动PyTorch 显存管理会更灵活。把双视角改成单视角显存占用几乎减半。降低vision_encoder.core_net_kwargs.feature_dim比如从 64 降到 32。这些办法按优先级尝试一般到第 4 步就能解决。注意修改 batch_size 会影响训练动态调小之后 lr 可以适当降一点或者多跑几个 epoch 让模型稳定下来。5.3 数据文件损坏后的修复思路训练过程中如果突然报KeyError或RuntimeError: Unable to synchronously read value很有可能是 HDF5 文件损坏或者读写冲突。常见情况有两种一是你一边训练一边往同一个 HDF5 里写数据导致文件锁冲突二是磁盘空间满HDF5 写入失败。前者可以换一个数据副本或者把数据只读打开后者需要清理磁盘空间后重新导出数据。还有一种隐蔽的问题是数据里混入了 NaN。低维状态或动作里有 NaN 时loss 会变成 NaN训练直接崩。排查方法是用 h5py 遍历所有 demo检查actions和关键 obs 是否有 NaNimport numpy as np import h5py with h5py.File(../datasets/lift_ph_image.hdf5, r) as f: for k in f[data].keys(): if not np.isfinite(f[fdata/{k}/actions][:]).all(): print(NaN in actions:, k)如果确认某个 demo 有 NaN写一个小脚本把它从 hdf5 里剔除或者单独屏蔽它而不是整个数据集重新下载。5.4 训练效果异常的排查顺序loss 正常下降但 rollout 成功率却很低或者 val loss 很高这种情况最容易让人抓狂。我整理的排查顺序是先确认是不是数据问题换一个已经训练好的 checkpoint 做 rollout如果成功率也低问题可能在数据本身或环境配置。检查 obs_keys 和动作空间是否匹配图像是否被错误地缩放到 0-255 再归一化、动作是否漏了某个维度。确认 GMM 是否开启如果演示动作分布是多峰的比如有的轨迹从左边抓、有的从右边抓MSE 回归会学一个平均动作成功率一定上不去。看训练日志中 loss 是否出现周期性尖刺尖刺往往代表某个难样本在反复捣乱可以尝试把该 demo 剔除。这一套排查下来大部分“训练玄学”都能被定位。我自己最常踩的坑是第一点——checkpoint 选错了评估用了一个早期保存的模型却以为当前模型效果差。6. 评估success rate 怎么算策略到底行不行6.1 rollout 评估的完整命令训练完成后评估是判断策略好坏的唯一标准。robomimic 提供run_trained_agent.py脚本可以对训练好的 agent 做环境 rolloutpython run_trained_agent.py \ --agent ../ckpt/lift_ph_vbc/trained_model.pt \ --config ../ckpt/lift_ph_vbc/config.json \ --episodes 50 \ --render这里trained_model.pt是训练结束后自动保存的最终模型--episodes 50表示连续跑 50 个测试回合--render开启可视化能看到机械臂在每个 episode 里的实际表现。如果想在无显示环境跑把--render去掉只输出数值结果。6.2 看懂评估输出评估结果会输出类似这样的内容Rollout: 1000, Success Rate: 0.84, Average Return: 0.84 rollout 0000: return1.00, success1 rollout 0001: return0.00, success0 ...Success Rate 就是成功回合占所有回合的比例。Lift 任务的成功标准是方块被机械臂抬升到目标高度以上这是一个明确的稀疏奖励所以Average Return在数值上通常就等于成功率成功得 1 分失败得 0 分。在视觉 BC 上我的经验是 Lift-PH 成功率至少要达到 0.7 以上才算“训练基本成功”。如果只有 0.3-0.5说明策略学到了一些行为但不稳定如果低于 0.2基本等于没学会需要回到前面的环节排查。6.3 失败模式与典型原因评估时不要只看最终数字要看具体的失败模式。通过--render观察我总结过 Lift-PH 常见的三种失败机械臂完全够不到方块大概率是视觉编码没学好或者仅靠图像时网络没能准确估计方块位置。这时可以试试加回低维状态或者延长训练。夹爪到达方块附近但不闭合可能是 GMM 输出的夹爪动作分布有问题或者演示数据里闭合夹爪的时机样本较少。抓起来但抬升高度不够常见于动作维度里抬升速度不够导致在回合结束前没有到达目标高度。可以检查演示数据里是否包含足够的抬升阶段轨迹。先纠正失败模式再追求成功率比盲目调参高效得多。比如“抓起来但抬升不够”你去调视觉编码器是没用的应该从动作层和演示轨迹分布入手。7. 从 Lift-PH 出发后续可以继续做的事7.1 把方案迁到 Can、Push 任务Lift-PH 跑通后robomimic 系列里还有 Can开罐、Push推物等任务它们的数据集结构和训练配置高度相似。迁移时只需要注意三件事一是任务的成功判定标准不同比如 Push 要求把物体推到目标区域需要看对应的 reward 定义二是动作空间可能不一样需要同步修改ac_dim三是 obs 里的视觉 key 可能命名不同先检查数据集结构再改配置。我在迁移到 Can 任务时遇到过一个问题Can 任务里手眼相机视角的作用比 Lift 更明显去掉eye_in_hand_image后成功率掉了不少所以建议这类接触操作任务尽量保留双视角。7.2 改进方向更强的骨干、更强的时序建模Vision BC 只是视觉模仿学习的起点。如果想把成功率推得更高一个方向是把单帧视觉编码升级为多帧观测让策略感知物体运动趋势另一个方向是引入 RNN 或 Transformer 结构比如 robomimic 里的 BC-RNN直接对时间序列建模在需要连续操作的任务上往往比单步 BC 更稳。图像增强在仿真里可能收益不大但如果未来要迁移到真实机器人建议跨场景验证增强策略让视觉编码器对光照和背景变化更鲁棒。我个人实际训练下来的体会是阶段 7 最值得投入时间的是数据构建和故障恢复这两个环节因为 Vision BC 的训练本身并不玄学很多失败其实在进入训练之前就已经注定了。把 HDF5 看明白、把 checkpoint 机制用好、把评估指标和失败模式对应起来这套流程无论之后跑什么模仿学习任务都能直接复用。最后再分享一个小技巧训练完成后把训练过程中保存的多个 checkpoint 都拿出来跑一遍评估不要只用最后一个trained_model.pt。有时候某个中间的 checkpoint 由于学习率衰减刚好落在更好的局部最优成功率反而比最终模型更高选它再继续微调能省下不少重新训练的时间。
robomimic 的 Lift-PH 复现系列终于写到阶段 7这是整个项目里视觉模仿学习最关键的一段Vision BC视觉行为克隆从数据构建、视觉编码、正式训练到故障恢复和最终评估的完整链路。前几个阶段如果已经把环境、数据集和低维 BC 跑通那么这一阶段就是把策略从“只看状态向量”升级到“真正看懂图像”难度跨了一个台阶但效果也最贴近真实机器人操作。这一篇我按照自己实际跑通 Lift-PH Vision BC 的过程整理重点不是贴文档而是回答那些文档里不写、但你跑起来一定会遇到的问题图像观测怎么从 HDF5 里取出来、ResNet 视觉编码器怎么和动作头拼在一起、训练中断了怎么优雅续上、成功率到底怎么算。适合正在复现 robomimic、想把视觉模仿学习真正跑起来的同学也适合已经跑通低维 BC、准备往视觉方向进阶的读者。1. 阶段全貌Vision BC 在 Lift-PH 中的定位与复现链路1.1 Lift-PH 是什么为什么阶段 7 要交给 Vision BCLift 是 robomimic 里最经典的仿真任务一个机械臂需要从桌面上抓起一个方块并抬升到目标高度。任务本身感知上不复杂但它是验证模仿学习算法的“最小可用环境”——动作空间小、状态变化直观、失败模式容易排查。后缀 PH 指的是 Proficient Human表示这批演示数据来自熟练人类操作者的遥操作记录而不是脚本或低质量随机策略。PH 数据的质量是后续一切训练的前提robomimic 官方把 Lift-PH 作为一个标准基准几乎所有算法BC、BC-RNN、IRL、CQL 等都在这个任务上有公开对比。阶段 7 之所以要引入 Vision BC是因为真实机器人场景里你不可能总是拿到完美的低维状态比如夹爪位置、物体坐标更多时候只有一个摄像头画面。Vision BC 直接从图像观测学习动作映射是打通“感知-控制”最简单直接的一步。它不需要强化学习那样的在线交互只需要离线演示数据训练完成后用 rollout 去验证策略好坏非常适合作为视觉模仿学习的第一个实践项目。1.2 从数据到部署的一条完整链路阶段 7 并不只是运行一行训练命令那么简单前面有数据构建后面有故障恢复和评估整条链路是确认演示数据Lift-PH 的 HDF5 文件里有哪些字段、图像是什么尺寸、动作空间是几维。构建训练数据集把 HDF5 中的原始观测转换成 robomimic Dataset 所需的观测字典配置 obs_keys、图像归一化、图像增强策略。设计视觉编码器搭建 ResNet 骨干、空间 softmax、与低维状态/动作的融合方式。正式训练写好 config 文件启动 train.py监控 action loss 和验证集表现。故障恢复训练中断后利用 checkpoint 续训处理显存不足和数据异常。评估与验证用 run_trained_agent.py 做 rollout计算成功率分析失败模式。这一步和前面低维 BC 最大的区别在于图像观测的引入让数据 pipeline 和网络结构都复杂了。图像数据体积大读取/预处理容易成为瓶颈视觉编码器参数量大显存和训练时间的压力也明显上升。所以阶段 7 的“数据构建”和“视觉编码”两个环节其实是整个 Vision BC 能否顺利训练的关键值得专门拆开讲。2. 数据构建把 HDF5 演示变成模型能吃的“视觉动作”样本2.1 先看清数据集结构别跳过这一步robomimic 的演示数据统一存放在 HDF5 文件里。Lift-PH 的官方路径一般是datasets/lift_ph_image.hdf5但我建议你用自己的数据目录避免和官方默认路径搞混。拿到数据集后第一件事不是急着训练而是用脚本检查它的内部结构import h5py dataset_path ../datasets/lift_ph_image.hdf5 with h5py.File(dataset_path, r) as f: print(root keys:, list(f.keys())) demo_keys list(f[data].keys()) print(num demos:, len(demo_keys)) demo f[data/demo_0] print(obs keys:, list(demo[obs].keys())) print(actions shape:, demo[actions].shape) print(image shape:, demo[obs][agentview_image].shape)这个检查非常关键。我遇到过不止一次数据集版本不对、obs key 命名不同导致后面训练时报 KeyError。Lift-PH 的agentview_image和eye_in_hand_image是常用的两个视觉观测shape 一般是(T, 3, H, W)通道在前或(T, H, W, 3)robomimic 会按 config 里设定的方式处理。动作维度 Lift 一般是 7 维即末端位置 3 维 四元数姿态 4 维有的数据版本会额外包含夹爪指令那就是 8 维或 10 维。你要在 check 时确认actions的维度然后在 config 里把ac_dim配成对应值。2.2 用 dataset_states_to_obs_dict_replay 完成观测转换robomimic 官方训练脚本里数据加载通常不直接操作 HDF5而是通过SequenceDataset配合一个 obs 字典转换函数完成。在 Vision BC 的配置里你需要明确告诉 robomimic哪些 key 来自 demo 的obs组哪些 key 是图像、哪些是低维状态。比较省事的做法是在 config 里配置data.obs_keys然后让 robomimic 内部完成转换。不过如果你的数据不是官方格式或者你想增加自定义特征就需要自己写一个转换函数def dataset_states_to_obs_dict_replay(demo): 将 hdf5 demo 中的原始 obs 转为 robomimic 训练所需的 obs 字典。 obs_dict { agentview_image: demo[obs][agentview_image], eye_in_hand_image: demo[obs][eye_in_hand_image], robot0_eef_pos: demo[obs][robot0_eef_pos], robot0_eef_quat: demo[obs][robot0_eef_quat], } return obs_dict然后传给 SequenceDataset 时指定这个函数。这样做的好处是你能完全控制哪些观测被送进网络。我在实际操作中习惯保留robot0_eef_pos和robot0_eef_quat这两个低维状态跟图像一起输入因为机械臂末端位姿对 Lift 这种需要精确到达物体的任务特别有用只靠图像会让策略学得更慢。2.3 图像像素预处理与归一化直接影响收敛速度很多人忽略图像归一化直接把 0-255 的 uint8 像素喂给 ResNet。这在 robomimic 里通常也能跑但收敛明显更慢而且 loss 容易出现奇怪的跳动。robomimic 提供data.image_obs_normalization之类的配置但更通用的做法是在 Dataset 的 obs 处理里对图像做 ImageNet 归一化transform torchvision.transforms.Normalize( mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] )注意归一化操作要在把 uint8 转成 float 并缩放到 [0,1] 之后进行。如果你用的是 robomimic 的ImageNormalize工具它会自动处理这些但自己写转换函数时一定要记得这一步。经历过一次效果对比后我的结论是归一化对 Vision BC 来说不是可选项而是必选项。它让 ResNet 的预训练权重如果用了发挥正常作用也避免不同光照条件下像素值偏移对训练造成干扰。2.4 数据质量自检清单数据是模仿学习的上限。Lift-PH 演示数据虽然叫“Proficient Human”但不代表每个 demo 都是完美可用的。我在训练前通常会做一套快速自检动作序列是否有长时间静止片段若有可能是遥操作故障这类 demo 会让策略学到“什么都不做”。图像是否出现全黑/全白帧这些帧会污染视觉编码器最好直接滤除。action 数值分布是否合理比如连续动作是否有 NaN末端位置是否突然跳变。所有 demo 的 obs key 是否一致不同 demo 缺少某个 key 是训练崩溃常见原因。建议写一个几行的检查脚本把每个 demo 的 actions 和图像帧做统计把异常的 demo id 打印出来。这一步虽然不产生“成果”但能在训练前拦截大量问题省掉后面排查的几小时。3. 视觉编码ResNet18、空间 softmax 与多视角融合3.1 为什么 Lift 场景必须看图像如果你已经跑过 Lift-PH 的低维 BC会发现只用robot0_eef_pos、robot0_eef_quat、物体位置这类状态也能学到不错的成功率。那 Vision BC 还有必要吗有而且价值体现在两点一是泛化低维状态在仿真里可以直接拿到但迁移到真实机器人时往往只能靠视觉二是信息密度图像里包含物体形状、位置、遮挡关系、夹爪开合状态这些隐含特征很难手工设计成低维向量。具体到 Lift 任务机械臂要先接近方块、调整夹爪方向、闭合夹爪、抬升。这个过程里“夹爪和物体的相对位置关系”是最核心的信息在图像里一目了然但在低维状态里需要精确的坐标变换才能表达。Vision BC 让网络自己从像素里提取这些关系省去了大量人工特征工程。3.2 视觉骨干网络搭建与参数选择robomimic 的 Vision BC 默认支持多种视觉骨干最常用的是 ResNet18。相比 ResNet50 或更重的网络ResNet18 在 Lift 这类简单视觉任务上完全够用训练更快显存占用更低对新手更友好。我的配置里一般这样设vision_encoder: { core_net_class: ResNet18Conv, core_net_kwargs: { input_channel: 3, feature_dim: 64, use_spatial_softmax: true, spatial_softmax_kwargs: { num_kp: 32, temperature: 1.0, learnable_temperature: true, output_size: 64 } } }这里use_spatial_softmax很关键。空间 softmax 会把卷积特征图上的每个位置转换成一个 2D 关键点坐标相当于让网络自己学会“关注图像里哪些位置”。对 Lift 来说这些关键点往往会落在方块边缘、夹爪指尖、目标区域附近可解释性很强。我建议在阶段 7 一定开启 spatial softmax它能让视觉特征更紧凑也更容易和后面的 MLP 层配合。3.3 视觉特征与低维状态特征的融合方式Vision BC 并不一定只用图像。前面我提到会保留robot0_eef_pos和robot0_eef_quat这时网络就需要处理两类输入视觉特征和低维状态。robomimic 的处理方式是把它们 concat 起来再送入后续的 MLP 层。具体来说图像经过 ResNet18 spatial softmax 输出一个feature_dim维的向量比如 64 维低维状态经过一个小的 MLP 编码成相同或不同维度的向量两者拼接后进入 actor 网络的隐层。这个设计的好处是视觉负责“看”状态负责“感知自身”两者互补。配置时注意actor_net.vision_encoder的输出维度要和低维状态编码后的维度对齐否则 concat 时会出问题。我在实际调试中的一个心得如果你发现网络训练正常但 rollout 时机械臂到处乱摸可以先试试只保留视觉输入去掉低维状态看成功率是否变化。这能帮你判断网络到底依赖哪路信息。如果去掉低维状态成功率骤降说明视觉编码还没学好如果成功率不变说明视觉已经掌握了核心信息低维状态只是辅助。3.4 图像增强的取舍别照搬分类任务的套路图像增强在视觉模仿学习里是个双刃剑。robomimic 支持随机裁剪、颜色抖动、翻转等增强方式。在 Lift 仿真环境里背景是固定的物体颜色也是固定的我实际测试下来随机裁剪有时反而会切掉机械臂末端导致策略困惑颜色抖动对收敛也没有明显帮助。我的建议是阶段 7 先关闭图像增强只做归一化。等 baseline 跑通之后再试着开一个轻微的随机裁剪比如裁剪后 resize 回原尺寸看看 val loss 和 rollout 成功率是否提升。不要在初期就引入复杂增强否则你会分不清是增强的问题还是网络结构的问题。仿真环境里固定的视觉分布本身就不是增强的用武之地增强真正发挥价值是在真实机器人数据跨场景迁移的时候。4. 正式训练config 逐项拆解与 loss 观测4.1 一份可用的 Vision BC 配置参考下面是我在 Lift-PH 上实际用过的精简配置去掉了与阶段 7 无关的冗余项保留核心参数{ experiment: { name: lift_ph_vbc, save: { enabled: true, frequency: 500 }, validate: { enabled: true, frequency: 10 } }, train: { num_epochs: 600, batch_size: 32, optimizer: { type: adamw, lr: 0.0003, weight_decay: 0.0001 }, output_dir: ../ckpt, checkpoint_frequency: 10, log_frequency: 100 }, algo: { share_encoder: false, actor_net: { ac_dim: 7, hidden_dim: 1024, layers: 3, vision_encoder: { core_net_class: ResNet18Conv, core_net_kwargs: { input_channel: 3, feature_dim: 64, use_spatial_softmax: true, spatial_softmax_kwargs: { num_kp: 32, temperature: 1.0, learnable_temperature: true, output_size: 64 } } } }, gmm: { enabled: true, num_modes: 5, min_std: 0.0001 } }, data: { dataset_path: ../datasets/lift_ph_image.hdf5, obs_keys: [agentview_image, eye_in_hand_image, robot0_eef_pos, robot0_eef_quat], image_obs_normalization: true, image_augmentation: false } }几个参数我说一下选择原因。batch_size32在单卡上处理双视角 ResNet18 是比较稳妥的值再大容易出现显存不足再小则训练不稳定。lr0.0003配 AdamW 是视觉主干网络常用的初始学习率如果 loss 震荡可以降到 0.0001。gmm.enabledtrue打开 GMM 输出头让网络预测动作的混合高斯分布这对模仿多模态演示比如同一种任务存在多种轨迹非常关键比直接回归动作效果好得多。4.2 训练启动与 loss 日志解读启动训练用 robomimic 的train.py把上面配置存成lift_ph_vision_bc.json然后执行python train.py \ --config configs/lift_ph_vision_bc.json \ --dataset ../datasets/lift_ph_image.hdf5训练日志会输出到output_dir下的lift_ph_vbc目录。重点关注两个指标train/action_loss和val/action_loss。前者是训练集上的动作预测误差后者是验证集上的误差。我观测到的一个典型过程是前 10 个 epochloss 从初始值快速下降val loss 也跟着降到 50-100 epoch 后train loss 继续缓慢下降val loss 开始波动甚至上升这时模型可能进入过拟合。robomimic 的 validate 设置每 10 个 epoch 评估一次你可以根据 val loss 的趋势决定是否提前终止而不是傻等到 600 epoch。如果显存和算力允许建议开 TensorBoard 或至少保存 csv 日志方便回看。4.3 显存占用与训练时长估算Vision BC 的显存压力比低维 BC 大得多尤其开了双视角图像。ResNet18 batch_size32 双视角在单张 12GB 显存卡上大概会占到 8-10GB属于比较紧的状态。如果显存不足优先做两件事减小 batch_size 到 16 或 8如果还不行就把eye_in_hand_image从 obs_keys 里去掉只留一个agentview_image。在 Lift 任务上手眼相机视角确实有帮助但少一个视角并不会让策略完全失效。训练时长方面600 epoch 在单张消费级显卡RTX 3080 级别上双视角 ResNet18 大约需要 1.5 到 3 小时。如果只是验证流程可以先设 50 epoch 跑通一遍确认日志、checkpoint、评估脚本都正常后再开长训练。我习惯的做法是先用 10 个 epoch “冒烟测试”检查代码和配置没有低级错误再正式跑长训练。5. 故障恢复中断续训、OOM 与脏数据急救手册5.1 用 checkpoint 回到断点别让训练从头再来训练跑了一半因为服务器重启、显存被占、或者手滑关掉终端导致中断是再常见不过的事。robomimic 默认会在output_dir下保存 checkpoint 文件命名类似checkpoint_epoch_50.pt、checkpoint_epoch_100.pt。恢复训练用--resume参数python train.py \ --config ../ckpt/lift_ph_vbc/config.json \ --resume ../ckpt/lift_ph_vbc/checkpoint_epoch_50.pt注意--resume会加载模型权重和优化器状态并继续从 checkpoint 记录的 epoch 开始训练。如果你只想加载模型权重但重新开始训练则不要用--resume而是手动设置train.start_epoch0并加载权重。这两个场景差别很大别搞混。我遇到过一次“续训后 loss 暴涨”的问题原因是续训时 config 里的 batch_size 被改了优化器状态和新的 batch 不匹配。恢复训练时最好保持和中断前完全一致的配置。5.2 显存不足的应急方案OOMOut of Memory是 Vision BC 训练中最常见的崩溃原因。如果你的日志里出现CUDA out of memory先别急着调网络结构按下面的顺序排查减小 batch_size从 32 到 16再到 8找到能跑稳的值。检查是否同时运行了多个占用显存的进程nvidia-smi看一下。如果板卡支持设置PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True再启动PyTorch 显存管理会更灵活。把双视角改成单视角显存占用几乎减半。降低vision_encoder.core_net_kwargs.feature_dim比如从 64 降到 32。这些办法按优先级尝试一般到第 4 步就能解决。注意修改 batch_size 会影响训练动态调小之后 lr 可以适当降一点或者多跑几个 epoch 让模型稳定下来。5.3 数据文件损坏后的修复思路训练过程中如果突然报KeyError或RuntimeError: Unable to synchronously read value很有可能是 HDF5 文件损坏或者读写冲突。常见情况有两种一是你一边训练一边往同一个 HDF5 里写数据导致文件锁冲突二是磁盘空间满HDF5 写入失败。前者可以换一个数据副本或者把数据只读打开后者需要清理磁盘空间后重新导出数据。还有一种隐蔽的问题是数据里混入了 NaN。低维状态或动作里有 NaN 时loss 会变成 NaN训练直接崩。排查方法是用 h5py 遍历所有 demo检查actions和关键 obs 是否有 NaNimport numpy as np import h5py with h5py.File(../datasets/lift_ph_image.hdf5, r) as f: for k in f[data].keys(): if not np.isfinite(f[fdata/{k}/actions][:]).all(): print(NaN in actions:, k)如果确认某个 demo 有 NaN写一个小脚本把它从 hdf5 里剔除或者单独屏蔽它而不是整个数据集重新下载。5.4 训练效果异常的排查顺序loss 正常下降但 rollout 成功率却很低或者 val loss 很高这种情况最容易让人抓狂。我整理的排查顺序是先确认是不是数据问题换一个已经训练好的 checkpoint 做 rollout如果成功率也低问题可能在数据本身或环境配置。检查 obs_keys 和动作空间是否匹配图像是否被错误地缩放到 0-255 再归一化、动作是否漏了某个维度。确认 GMM 是否开启如果演示动作分布是多峰的比如有的轨迹从左边抓、有的从右边抓MSE 回归会学一个平均动作成功率一定上不去。看训练日志中 loss 是否出现周期性尖刺尖刺往往代表某个难样本在反复捣乱可以尝试把该 demo 剔除。这一套排查下来大部分“训练玄学”都能被定位。我自己最常踩的坑是第一点——checkpoint 选错了评估用了一个早期保存的模型却以为当前模型效果差。6. 评估success rate 怎么算策略到底行不行6.1 rollout 评估的完整命令训练完成后评估是判断策略好坏的唯一标准。robomimic 提供run_trained_agent.py脚本可以对训练好的 agent 做环境 rolloutpython run_trained_agent.py \ --agent ../ckpt/lift_ph_vbc/trained_model.pt \ --config ../ckpt/lift_ph_vbc/config.json \ --episodes 50 \ --render这里trained_model.pt是训练结束后自动保存的最终模型--episodes 50表示连续跑 50 个测试回合--render开启可视化能看到机械臂在每个 episode 里的实际表现。如果想在无显示环境跑把--render去掉只输出数值结果。6.2 看懂评估输出评估结果会输出类似这样的内容Rollout: 1000, Success Rate: 0.84, Average Return: 0.84 rollout 0000: return1.00, success1 rollout 0001: return0.00, success0 ...Success Rate 就是成功回合占所有回合的比例。Lift 任务的成功标准是方块被机械臂抬升到目标高度以上这是一个明确的稀疏奖励所以Average Return在数值上通常就等于成功率成功得 1 分失败得 0 分。在视觉 BC 上我的经验是 Lift-PH 成功率至少要达到 0.7 以上才算“训练基本成功”。如果只有 0.3-0.5说明策略学到了一些行为但不稳定如果低于 0.2基本等于没学会需要回到前面的环节排查。6.3 失败模式与典型原因评估时不要只看最终数字要看具体的失败模式。通过--render观察我总结过 Lift-PH 常见的三种失败机械臂完全够不到方块大概率是视觉编码没学好或者仅靠图像时网络没能准确估计方块位置。这时可以试试加回低维状态或者延长训练。夹爪到达方块附近但不闭合可能是 GMM 输出的夹爪动作分布有问题或者演示数据里闭合夹爪的时机样本较少。抓起来但抬升高度不够常见于动作维度里抬升速度不够导致在回合结束前没有到达目标高度。可以检查演示数据里是否包含足够的抬升阶段轨迹。先纠正失败模式再追求成功率比盲目调参高效得多。比如“抓起来但抬升不够”你去调视觉编码器是没用的应该从动作层和演示轨迹分布入手。7. 从 Lift-PH 出发后续可以继续做的事7.1 把方案迁到 Can、Push 任务Lift-PH 跑通后robomimic 系列里还有 Can开罐、Push推物等任务它们的数据集结构和训练配置高度相似。迁移时只需要注意三件事一是任务的成功判定标准不同比如 Push 要求把物体推到目标区域需要看对应的 reward 定义二是动作空间可能不一样需要同步修改ac_dim三是 obs 里的视觉 key 可能命名不同先检查数据集结构再改配置。我在迁移到 Can 任务时遇到过一个问题Can 任务里手眼相机视角的作用比 Lift 更明显去掉eye_in_hand_image后成功率掉了不少所以建议这类接触操作任务尽量保留双视角。7.2 改进方向更强的骨干、更强的时序建模Vision BC 只是视觉模仿学习的起点。如果想把成功率推得更高一个方向是把单帧视觉编码升级为多帧观测让策略感知物体运动趋势另一个方向是引入 RNN 或 Transformer 结构比如 robomimic 里的 BC-RNN直接对时间序列建模在需要连续操作的任务上往往比单步 BC 更稳。图像增强在仿真里可能收益不大但如果未来要迁移到真实机器人建议跨场景验证增强策略让视觉编码器对光照和背景变化更鲁棒。我个人实际训练下来的体会是阶段 7 最值得投入时间的是数据构建和故障恢复这两个环节因为 Vision BC 的训练本身并不玄学很多失败其实在进入训练之前就已经注定了。把 HDF5 看明白、把 checkpoint 机制用好、把评估指标和失败模式对应起来这套流程无论之后跑什么模仿学习任务都能直接复用。最后再分享一个小技巧训练完成后把训练过程中保存的多个 checkpoint 都拿出来跑一遍评估不要只用最后一个trained_model.pt。有时候某个中间的 checkpoint 由于学习率衰减刚好落在更好的局部最优成功率反而比最终模型更高选它再继续微调能省下不少重新训练的时间。
