这个教程来自这个英伟达网址
教程的目标是提供用于部署 VLA 模型的优化量化和推理方法,以及针对新机器人、任务和环境的参考微调流程。在一个自包含的仿真环境中,结合场景生成和领域随机化(MimicGen)对性能和准确性进行严格验证。未来阶段将包括与 Isaac Lab 和 ROS2 的 sim2real 集成、对 CrossFormer 等相关模型的研究,以及针对实时性能的神经网络结构优化。
✅ 针对 VLA 模型的量化和推理优化
✅ 原始 OpenVLA-7B 权重的准确性验证
✅ 基于合成数据生成的参考微调工作流程
✅ 在 Jetson AGX Orin 上使用 LoRA 进行设备端训练,以及在 A100/H100 实例上进行完全微调
✅ 在示例积木堆叠任务中通过领域随机化达到 85% 的准确率
✅ 提供用于复现结果的示例数据集和测试模型
1. 量化
已在 NanoLLM 的流式 VLM 管道中添加对 OpenVLA 的支持,使用 MLC 在 INT4/FP8 量化下运行,并利用 TensorRT 以 FP16 精度执行视觉编码器。
首先,我们将在 BridgeData V2 数据集(来自 Open X-Embodiment 集合中权重最高的数据集之一)上测试模型。该数据集也是模型的训练集,用于在部署时验证量化和推理是否正确。以下截图摘自其官网,展示了数据集的样貌:
下面的命令会启动容器、下载数据集和模型(如有必要)、在首次运行时进行量化,并使用归一化均方误差(NRMSE)来衡量模型预测动作值与数据集真实值之间的准确度,以消除动作空间各个维度范围不一致带来的偏差。我们在 HuggingFace Hub 上抽取了原始 Bridge 数据集的 100 个 episode 子集,因此无需为这些测试下载完整的约 400 GB 数据集。
INT4:
jetson-containers run $(autotag nano_llm) \
python3 -m nano_llm.vision.vla --api mlc \
--model openvla/openvla-7b \
--quantization q4f16_ft \
--dataset dusty-nv/bridge_orig_ep100 \
--dataset-type rlds \
--max-episodes 10 \
--save-stats /data/benchmarks/openvla_bridge_int4.json
FP8
jetson-containers run $(autotag nano_llm) \
python3 -m nano_llm.vision.vla --api mlc \
--model openvla/openvla-7b \
--quantization q8f16_ft \
--dataset dusty-nv/bridge_orig_ep100 \
--dataset-type rlds \
--max-episodes 10 \
--save-stats /data/benchmarks/openvla_bridge_fp8.json
FP16
jetson-containers run $(autotag nano_llm) \
python3 -m nano_llm.vision.vla --api hf \
--model openvla/openvla-7b \
--dataset dusty-nv/bridge_orig_ep100 \
--dataset-type rlds \
--max-episodes 10 \
--save-stats /data/benchmarks/openvla_bridge_fp16.json
使用
--save-stats
可以将逐帧指标和平均值保存到文件;使用
--max-episodes
可以限制处理的 episode 数量。
测试环境:Jetson AGX Orin 64GB,JetPack 6
量化精度对比:后续微调模型在 INT4 上的准确度将与 FP8/FP16 保持一致
每一帧都会打印模型预测的 7 维动作向量及其真实值,以及该帧的准确率、延迟和帧率。波浪号(~)后的数字为迄今为止各项指标的平均值,命令结束时最后一行报告的是整个数据集处理完毕后的总体均值。
# INT4
step 355 [-0.02692 0.00776 -0.00299 0.08160 0.07292 0.04791 0.99608] accuracy 0.8466 ~0.9017 time=336.2 ms fps=2.96 ~2.97
gt 355 [-0.02387 0.00760 -0.00318 0.15965 0.07707 0.03281 1.00000]
# FP8
step 355 [-0.02392 0.00767 -0.00310 0.08160 0.07692 0.03217 0.99608] accuracy 0.9982 ~0.9523 time=469.7 ms fps=2.13 ~2.12
gt 355 [-0.02387 0.00760 -0.00318 0.15965 0.07707 0.03281 1.00000]
# FP16
step 355 [-0.02392 0.00767 -0.00310 0.08160 0.07692 0.03217 0.99608] accuracy 0.9982 ~0.9531 time=842.3 ms fps=1.19 ~1.18
gt 355 [-0.02387 0.00760 -0.00318 0.15965 0.07707 0.03281 1.00000]
如上所述,Bridge 数据集已包含在训练集中,下面我们还会在未参与训练、具有更大场景变化的新数据上再次运行测试。此工具同样支持加载 Open X-Embodiment 中的 RLDS/TFDS 格式数据集,以及 Robomimic/MimicGen 的 HDF5 格式数据。你也可以参考下面的代码示例,利用公开的 API 自行创建 agent 和测试脚本。
2. 推理API
在容器中使用 NanoLLM 库对摄像头流运行 VLA 推理的代码非常简单:
from nano_llm import NanoLLM
from nano_llm.plugins import VideoSource
# 加载视觉-语言-动作模型(VLA),使用 q4f16_ft 量化方式
model = NanoLLM.from_pretrained(model, quantization='q4f16_ft')
# 打开摄像头(或视频流),在 CUDA 流 0 上进行处理
camera = VideoSource(video_source, cuda_stream=0)
assert(model.vla) # 确保加载的是 VLA 模型
while True:
# 捕获一帧图像,返回格式可以是 cudaImage、NumPy 数组或 GPU 上的 Torch 张量
image = camera.capture()
if image is None: # 如果因超时没捕获到图像,就继续重试
continue
# 基于当前图像和指令,预测动作向量
# 返回一个长度等于自由度数(dof)的数组或张量
# 对于 OpenVLA,顺序是 (Δx, Δy, Δz, Δroll, Δpitch, Δyaw, gripper)
actions = model.vla.predict_action(
image,
instruction="pick up the nearest object", # “捡起最近的物体”
action_space="normalized", # 返回已归一化的动作值
return_tensors='np', # 用 NumPy 数组格式返回
)
# 接着就可以把 actions 发给机器人控制器或 IK 解算器去执行啦
...
quantization='q4f16_ft'
:在推理时使用 INT4 + FP16 混合量化,并进行过微调,以在嵌入式设备(如 Jetson)上获得最佳性能和精度平衡。
VLA 模型也同样被 Agent Studio支持,该平台还集成了仿真组件。
3. 在线验证
鉴于该任务领域的挑战性、动态反馈回路以及仿真/训练/推理的计算需求,使用 VLA 进行语言引导的灵巧操作相比于仅使用 LLM 和 VLM,复杂度显著提升。
要从在 token 级别预测 logits,到在一系列连续帧中不断输出足够正确的动作以形成有用行为,就必须在训练和推理流程的每个阶段交叉检验输出并衡量准确度,以便在出现性能回退时能够定位原因。
与典型的监督学习应用不同,最终任务完成度和成功率并不是从不考虑时间序列和物理交互反馈及复合误差的静态预录数据集中测量的——它们需要在线验证,无论是在仿真中还是在真实测试中。
弥合仿真到现实的鸿沟:使用 NVIDIA Isaac Lab 训练 Spot 四足机运动
在训练期间,会根据真实动作标签(即有多少动作 token 被准确预测)来测量 token 分类准确率,模型通过最小化该损失来优化(这与 LLM 的常规做法一致)。
同时,也会根据还原出的实值输出与真实值的 L1 误差来衡量连续域动作准确率。由于 token 分类不对“更接近目标”提供额外奖励,连续动作准确率通常会略高于分类准确率。
实际上,为了在类似环境中成功完成任务,这两项准确率都应保持在 95%以上。
为了达到如此高的准确度,相关工作和研究中通常会故意让模型在相同的 90 万个 episode 上训练多达 30 个 epoch(OpenVLA),以实现“过拟合”。Transformer 擅长通过少量示例记忆特定知识,但也对过拟合和遗忘先前学到的信息非常敏感。因此,LLM 在常规微调中通常只训练少数几个 epoch,以保留其零样本能力和对分布外输入的泛化能力。
在本项目的微调阶段,我们评估了不同训练 episode 数量与重复数据 epoch 数量对模型准确率和任务成功率的影响。
直到推理阶段才会测量实际的任务成功率,此时模型要么连入仿真器,要么在相似条件下进行一系列耗时的物理实验。
我们将 MimicGen 直接集成到 OpenVLA 的训练脚本中,以源源不断地产生全新的数据,但在模型处理了大量 episode 后,遇到了梯度不稳定的问题。
4. 用MimicGen仿真
MimicGen 利用场景图信息和任务/子任务的元数据(指出环境中哪些物体是当前子任务的目标),仅凭 10 个远程示教样本,就能生成随机化的训练回合——将原始示教轨迹插值到新的、随机的位置和姿态中。
这样即可产生大量独特的训练数据,提高模型的鲁棒性,而无需投入大量人工来教机器人新技能和行为。
MimicGen 构建在 Robomimic 和 Robosuite 仿真器之上,可与 VLA 一同在 Jetson 上无头运行,极大简化了可复现性实验的环境搭建。RoboCasa 项目基于 MimicGen,目前正与 NVIDIA Omniverse 和 OSMO 集成。未来我们还计划引入 Isaac Lab,以获得更好的可扩展性、更精准的物理模拟和更逼真的渲染效果。
MimicGen 包含 12 种任务,例如块状物堆叠、抓取与放置、装配以及厨房场景。每种任务都有随着学习进展而难度递增的多个变体,这为对比“人工设计”的训练顺序与 OpenVLA 在 Open X‑Embodiment 中采用的纯随机化顺序,提供了很好的研究机会。
在本教程阶段,我们将聚焦“块状物堆叠”任务,以了解要达到论文中 >75–80% 的成功率,所需的训练需求和运行时性能指标。
这些分析将为在更多行为模式和更复杂的场景(如与 MimicGen 环境差异较大的部署环境)中扩展模型提供指导——正如原始 OpenVLA 权重在那些场景中完全没有成功记录所显示的那样。
4.1 数据生成
我们基于一份针对 aarch64+igpu 做了补丁的代码分支,构建了用于 Jetson 的 MimicGen 容器,并增加了一些功能增强:
为相关任务生成带有随机变体的自然语言标签
对物体的颜色/材质做额外的领域随机化
(对应的环境版本命名为Stack_D2
、Stack_D3
和Stack_D4
)
在训练 OpenVLA 时,会将图像和标签保存到磁盘;而后续的推理则在在线仿真中完成,以测量任务成功率。为此,我们将 MimicGen 与 Agent Studio 集成,可以交互式地测试模型,并快速插入诸如 ASR(自动语音识别)等组件,用以语音指令控制机器人。
在这个 OpenVLA 分支中,已初步支持将 MimicGen 直接集成到训练流程,实现实时仿真与验证,以及在无限新回合中训练而不重复同一批数据(epoch)。我们发现 LoRA 微调后期出现梯度突增的问题,建议可尝试降低学习率,或在完整微调脚本中同样集成 MimicGen,并使用 FDSP 在 dGPU 上增大批量大小来稳健训练。
下面的命令会生成指定数量的训练回合,保存为 Robomimic HDF5 格式。我们已在 HuggingFace Hub 上提供了包含 1000 和 2500 回合的渲染数据集。OpenVLA 建议仅需 10–150 回合即可完成高效微调与适应,其他类似场景中表现可能相近,但我们发现在 MimicGen 环境下远不足以达到预期效果。
jetson-containers run $(autotag nano_llm) \
python3 -m mimicgen.generate \
--tasks Stack_D4 \
--episodes 100 \
--output /data/datasets/mimicgen \
--cameras agentview \
--camera-width 224 \
--camera-height 224
4.3 RLDS 转换
OpenVLA 使用基于 TFDS 的 RLDS 格式数据集,所以我们提供了一个从 HDF5 到 RLDS 的转换工具。对于大量回合的数据,这一步也可能非常耗时,这正是我们希望在训练时在线运行 MimicGen 并直接将其集成到 OpenVLA 的原因之一。除非你需要生成不同的数据,否则可以跳过这一步,直接使用我们已经上传为 RLDS 格式的 MimicGen 数据集。
jetson-containers run $(autotag nano_llm) \
python3 -m nano_llm.datasets \
--dataset /data/datasets/mimicgen/demo_src_stack_task_D4/demo.hdf5 \
--dataset-type mimicgen \
--convert rlds \
--remap-keys agentview:image \
--output /data/datasets/mimicgen/rlds/stack_d4_ep2500
运行转换后,会在指定的输出目录下生成一组 TFRecord 文件,这些文件能够被 OpenVLA 的训练脚本直接加载使用。
5. 微调
本项目的主要目标之一是量化让模型适应不同机器人和任务所需的训练工作量。
我们的开发流程主要包括在 Jetson AGX Orin 64GB 上运行测试性 LoRA 微调并本地调试问题,当结果令人满意时,再在 Brev.dev、Vast.ai 和 RunPod 等提供商的 A100/H100 云实例上,使用 FDSP 完整地进行微调。我们也尝试过在两台 Jetson AGX Orin 上用 FDSP 做全量微调,但所有设置均因显存不足而失败。我们已将训练好的测试模型上传到 HuggingFace,供大家在下文的推理+仿真环境中试用。
以下是我们曾使用的训练 GPU 配置及相应的最大化显存利用的 batch size:
在 Jetson AGX Orin 和两块 A100 上训练至收敛大约需要 24–36 小时,具体时长取决于数据量和训练轮数(epoch)。为了尽量避免前文提到的过拟合,我们在完整微调时将 epoch 控制在 5 次以内,而通过增加训练回合数来提升模型表现。
具体过程可以看OpenVLA的官方repo
6. 推理和仿真
为了衡量模型在完成任务时的真实表现,我们在 Agent Studio 中启动了一个与 VLA 模型相连的 MimicGen 环境。它通过检查仿真发放的奖励来统计成功回合数——这些奖励模型本身并不使用,仅用以标记任务何时完成。我们设置了 200 帧的执行上限,超过该帧数则判定为失败。
jetson-containers run $(autotag nano_llm) \
python3 -m nano_llm.studio --load OpenVLA-MimicGen-INT4
jetson-containers run $(autotag nano_llm) \
python3 -m nano_llm.studio --load OpenVLA-MimicGen-FP8
jetson-containers run $(autotag nano_llm) \
python3 -m nano_llm.studio --load OpenVLA-MimicGen-FP16
要开始基准测试,只需将 AutoPrompt 节点的输出连接到 OpenVLA 节点。默认情况下它会无限运行——我们各做了 100 回合测试,由于仿真与模型推理同步运作,整个过程可能需要数小时(后续实验会尝试基于多步累积动作进行训练,并缩小模型规模以提升性能)。
这里每个方框都是一个「模块」,线条表示数据流向。
WebAudioIn
- 功能:从浏览器或麦克风抓取实时音频流。
- 输出:raw audio。
VADFilter
全称 Voice Activity Detector,检测「有人在说话」的时段,自动切掉静音与噪声。
输入:raw audio;输出:audio segments + interrupt 信号(暂中断)。
WhisperASR
用 OpenAI Whisper 做自动语音识别,输出 partial(中间结果)和 final(完整文字)。
作用:把语音转换成文字命令。
TextOverlay → VideoOutput
TextOverlay:把识别出来的文字直接叠加到视频帧上(右侧仿真窗口里看到的绿色提示)。
VideoOutput:负责把最终合成的图像以 ~5–6 FPS 推送到界面右侧。
UserPrompt
用户也可以手工在这里输入文字(比如在测试时直接打“把红块放到大块上”)。
AutoPrompt
把 WhisperASR 的结果、UserPrompt,以及 MimicGen 提供的上下文(见下)——综合出一个“List”格式的、给 LLM 的 Prompt 列表。
MimicGen
这里是视觉+语言的示范合成器:
image:当前仿真画面截图;
text:可以是示范轨迹的文字描述或环境标注;
它会在内部根据示例自动生成更多“虚拟示范”,增强 LLM/策略的少样本学习能力。
openvla-7b+stac
把上游的 Prompt 列表喂给一个 7B 参数量的 VLA(Vision-Language-Action)模型,输出一系列 actions(机器人动作指令)。
“2.8 FPS” 表示这个模型的推理速度。
这些 actions 再回流到 VideoOutput,驱动仿真里的机械臂做实际操作。
此次测试使用的是在 2500 回合数据上微调 4 个 epoch 的模型。尽管任务相对简单,该模型依然证明能够达到约 85% 的预期成功率。量化对性能的提升几乎呈线性增长,对准确率的影响不到 1%。
在评估了那些使用较少回合但训练更多轮次的微调模型的成功率后,我们清楚地看到了增大数据集规模所带来的影响:
这并不意味着使用 2500 回合训练的模型就没有过拟合——事实上,它已经学会了总是先堆放小方块,基本忽略了指令。要解决这个问题,就需要在训练中引入更多任务多样性,而现在我们已有了这样的数据流程。
与此同时,我们也希望将研究扩展到现实世界场景,而不仅仅局限于仿真环境,以验证模型在真实操作中的可行性。
在操控这一极具挑战的问题域中,要达到通用模型所需的变化量和数据量,不禁让人思考:是采用完全随机化的任务学习,还是在训练进程中随着复杂度递增、以更有针对性和精心设计的方式进行指导,更能高效地培养出具备广泛适应能力的模型?
7. 总结
搞了半天,应该去 Agent Studio 中找我想要的,如何启动一个与 VLA 模型相连的 MimicGen 环境(只有mimicgen?能否使用其他VLA?)。