全网首发! Nvidia Jetson Thor 128GB DK 刷机与测评(六)常用功能测评 pi0、pi0.5 环境部署
这篇博客是上一篇博客的续集,所有测试和评测均基于第一篇刷机博客的环境上完成的测试。此篇博客对 VLA 领域中最负盛名的模型 pi0
和 pi0.5
在 Jetson Thor 上进行了部署和测试。Thor 系列的其他博客如下:
- 全网首发! Nvidia Jetson Thor 128GB DK 刷机与测评(一)刷机与 OpenCV-CUDA、pytorch CUDA13.0+ 使用
- 全网首发! Nvidia Jetson Thor 128GB DK 刷机与测评(二)常用功能测评 Ollama、ChatTTS、Yolov13
- 全网首发! Nvidia Jetson Thor 128GB DK 刷机与测评(三)常用功能测评 DeepAnything 系列
- 全网首发! Nvidia Jetson Thor 128GB DK 刷机与测评(四)常用功能测评 - 目标跟踪系列 DeepSort、Boxmot、ByteTrack等
- 全网首发! Nvidia Jetson Thor 128GB DK 刷机与测评(五)常用功能测评 - RealtimeSTT 音频转文本 & 同声传译
🏫 Contributors
该工程由 中国科学院自动化研究所
提供硬件与测试场地,同时以下人员在该项目的开发中做出了巨大贡献,在此感谢他们的付出:
WenJiang Xu 徐文江,PengFei Yi 易鹏飞,JingKai Xu 徐靖凯,XingYu Wang 王行宇,YaNan Hao 郝亚楠,YuWei Wang 王雨薇
如果你对我们在 Jetson 设备上的其他工作感兴趣,可以访问我们的 Github 主页:
- https://github.com/GaohaoZhou-ops?tab=repositories
🍕 pi0 & pi0.5
如果你是做具身智能和 VLA,那么对这两个模型肯定不陌生,这两个模型可以说是目前所有 VLA 模型的参照物,也是必须掌握的模型。相关论文精读连接如下:
- VLA 论文精读(二十三)π0: A Vision-Language-Action Flow Model for General Robot Control
- VLA 论文精读(十八)π0.5: a Vision-Language-Action Model with Open-World Generalization
🚨 [必读] 前倾提要
在正式开始之前,一下几点是你必须了解的,每一条都很重要不能跳过。
- 由于 Jetson Thor、CUDA 13、Jax、pi0.5 都是比较新的架构和模型,在部署方面这些项目作者可能会有所改动,如果你不是着急使用话可以再等等官方进行适配,我们会持续跟踪这些作者的动态并在第一时间进行更新;
- 这篇博客的部署方法我们还没有进行充分测试,仅停留在实验阶段,不建议直接用于生产环境中;
- 我们将这篇博客中涉及到的相关资源放在了网盘中,包括 打包好的 Docker镜像 ,该镜像内置了
pi05_droid
两个模型; - 由于当前 jax、cuda、torch 之间的关联存在问题,因此推理效率不高,我们会 在所有优化都完成后上传一份新的镜像,如果你有能力解决也欢迎在评论区进行留言;
有关 Jax
和 CUDA 13
的适配,根据 Issue 上讨论,Jax 团队正在积极适配 CUDA 13,我们实验室也会尝试去解决这方面的问题:
- Issue 3118: https://github.com/jax-ml/jax/discussions/31118#discussioncomment-14142150
博客的剩余部分将介绍如何直接使用镜像以及这个镜像是如何构建的。
网盘链接:
通过网盘分享的文件:OpenPi
链接: https://pan.baidu.com/s/1mCWEdmJC2DMKtzlL4YD9Rg?pwd=iedf 提取码: iedf
--来自百度网盘超级会员v5的分享
💻 确认软硬件
为了避免冲突或歧义,首先需要确认你的软硬件是否和我的一致:
- Device:Nvidia Jetson Thor 128 GB DK;
- OS:Ubuntu 24.04;
- Jetpack:7.0;
- CUDA:V13.0.48;
- GPU Driver Version: 580.00;
Nvidia-smi | Jetpack Realease |
---|---|
![]() | ![]() |
☁️ 方式一:直接使用镜像
你需要将网盘中下载下面两个文件:
openpi_ws.zip
(867 MB):包含了所有源码;openpi.tar
(36 GB):完整镜像但不包含源码,注意不要下载成openpi.zip
;
Step1. 加载镜像
将 openpi.tar
文件下载下来之后使用下面的命令加载镜像:
$ sudo docker load -i openpi.tar
$ sudo docker images
Step2. 创建工作空间
将 openpi_ws.zip
文件解压到桌面:
$ cd Desktop
$ unzip openpi_ws.zip
Step3. 启动容器
使用下面的命令启动容器:
$ sudo docker run -dit \--net=host \--runtime nvidia \--privileged \--ipc=host \--ulimit memlock=-1 \--name openpi \--ulimit stack=67108864 \-v $(pwd):/workspace \openpi \bash$ sudo docker ps -a
- 启动容器后的推理部分两种方法一致,直接翻到最下面查看即可;
- 该镜像中包含了一个
pi05_droid
模型,如果你想要使用其他模型可以修改openpi/src/downloads.py
文件然后执行命令python src/downloads.py
下载,该文件是我们自己写的一个简单脚本,并非官方内容;
🐳 方式二:逐步构建镜像
Step1. 拉取基础镜像
这里使用 Nvidia 官方提供的 pytorch 镜像 nvcr.io/nvidia/pytorch:25.08-py3
作为底座,使用下面的命令拉取:
$ sudo docker pull nvcr.io/nvidia/pytorch:25.08-py3
$ sudo docker images
Step2. 创建工作空间并拉取源码
创建工作空间
在宿主机上创建一个工作空间用来在启动容器时映射,这里假设为 Desktop/openpi_ws
:
$ cd Desktop
$ mkdir openpi_ws
拉取 openpi 源码
在 openpi_ws
工作空间中拉取源码:
$ cd Desktop/openpi_ws
$ git clone https://github.com/Physical-Intelligence/openpi.git
创建 requirements.txt 文件
拉取完后进入到这个目录创建一个 requirements.txt
的文件并将下面的内容写入,此处不使用官方推荐的 uv 构建环境,因为存在太多冲突:
augmax==0.3.4
dm-tree==0.1.8
einops==0.8.0
equinox==0.11.8
flatbuffers==24.3.25
flax==0.10.2
fsspec[gcs]==2024.6.0
gym-aloha==0.1.1
imageio==2.36.1
# jax[cuda12]==0.5.3
jax[cuda13]==0.5.3
jaxtyping==0.2.36
# lerobot
ml_collections==1.0.0
# numpy==1.22.4<2.0.0
numpy==1.26.4
numpydantic==1.6.6
opencv-python==4.10.0.84
# openpi-client
orbax-checkpoint==0.11.13
pillow==11.0.0
sentencepiece==0.2.0
# torch==2.7.1
tqdm-loggable==0.2
typing-extensions==4.12.2
tyro==0.9.5
wandb==0.19.1
filelock==3.16.1
beartype==0.19.0
treescope==0.1.7
transformers==4.53.2
rich==14.0.0
polars==1.25
dm_control==1.0.14
创建 downloads.py 文件
为了能拉在推理之前取模型,在 openpi_ws/openpi/src
目录下新建一个 downloads.py
的文件并将下面的内容写入:
from openpi.training import config
from openpi.policies import policy_config
from openpi.shared import downloadmodel_name = "pi05_droid"
model_link = "gs://openpi-assets/checkpoints/pi05_droid"config = config.get_config(model_name)
checkpoint_dir = download.maybe_download(model_link)policy = policy_config.create_trained_policy(config, checkpoint_dir)
拉取 lerobot 源码:
在 openpi_ws
工作空间中拉取源码并切换到正确的分支:
$ $ cd Desktop/openpi_ws
$ git clone https://github.com/huggingface/lerobot.git
$ cd lerobot
$ git checkout 0cf864870cf29f4738d3ade893e6fd13fbd7cdb5
拉取 mujoco 轮子
前往 Release 页下载 2.3.7
轮子,或者用下面的命令下载:
$ cd Desktop/openpi_ws
$ wget https://github.com/google-deepmind/mujoco/releases/download/2.3.7/mujoco-2.3.7-linux-aarch64.tar.gz
$ tar -zxvf mujoco-2.3.7-linux-aarch64.tar.gz
完成后你的文件结构应该如下所示:
$ ll Desktop/openpi_ws
total 20K
drwxrwxr-x 5 thor thor 4.0K Sep 21 16:39 ./
drwxr-xr-x 9 thor thor 4.0K Sep 21 16:39 ../
drwxrwxr-x 12 thor thor 4.0K Sep 21 13:04 lerobot/
drwxr-sr-x 8 root 1005 4.0K Jul 20 2023 mujoco-2.3.7/
drwxrwxr-x 11 thor thor 4.0K Sep 21 12:08 openpi/
Step3. 运行容器
在宿主机上进入到刚才创建的工作空间中后使用下面的命令运行容器:
进入容器
$ cd Desktop/openpi_ws$ sudo docker run -dit \--net=host \--runtime nvidia \--privileged \--ipc=host \--ulimit memlock=-1 \--name openpi \--ulimit stack=67108864 \-v $(pwd):/workspace \nvcr.io/nvidia/pytorch:25.08-py3 \bash
然后进入容器:
$ sudo docker exec -it openpi /bin/bash
安装依赖库
更新 apt 源:
$ apt-get update
安装以下依赖库:
$ apt-get install libx11-dev libxext-dev libxrender-dev libxtst-dev libxt-dev
$ apt-get install libgl1-mesa-dev libxrandr-dev doxygen clang libopenexr-dev
$ apt-get install libxinerama-dev libxcursor-dev
$ apt-get install zip tmux tree
编辑 ~/.bashrc 文件
在 ~/.bashrc
文件末尾添加以下内容,将 CUDA 和 mujoco 目录添加进去:
export MUJOCO_PATH="/workspace/mujoco-2.3.7"
export MUJOCO_PLUGIN_PATH="/workspace/mujoco-2.3.7/bin/mujoco_plugin"
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-13.0/lib64
export PATH=$PATH:/usr/local/cuda-13.0/bin
export CUDA_HOME=$CUDA_HOME:/usr/local/cuda-13.0
export HF_ENDPOINT="https://hf-mirror.com"
保存完后 source 一下:
$ source ~/.bashrc
安装 mujoco 2.3.7
在容器中的工作空间 workspace
内使用下面的命令安装:
$ cd /workspace
$ pip install mujoco==2.3.7 -i https://pypi.tuna.tsinghua.edu.cn/simple
安装 openpi 依赖的第三方库
在容器中的工作空间 workspace/openpi
内使用下面的命令安装:
$ cd /workspace/openpi
$ pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
安装 openpi-client
进入到 /workspace/openpi/packages/openpi-client
目录下安装:
$ cd /workspace/openpi/packages/openpi-client
$ pip install -e . --no-deps
安装 lerobot
进入到 /workspace/lerobot/
目录下安装:
$ cd /workspace/lerobot/
$ pip install -e . --no-deps
安装其他第三方库
$ pip install datasets jsonlines draccus av websockets pydantic==2.10.6 -i https://pypi.tuna.tsinghua.edu.cn/simple
下面的报错可以忽略:
安装 openpi
进入到 /workspace/openpi
目录下执行以下命令:
$ cd /workspace/openpi
$ pip install -e . --no-deps
Step4. 下载模型
下载模型方面你可以选择使用我们上面添加的那个脚本进行下载,也可以从网盘中下载一个模型然后拷贝进容器并解压到 /root/.cache/openpi/openpi-assets/checkpoints
目录下,这里以脚本形式下载 pi05_driod
模型为例:
$ cd /workspace/openpi/
$ python src/downloads.py
【注意】:此处就出现了上面提到的 Jax 与 CUDA 的问题,我们实验室会尽可能去解决并在第一时间更新到博客中。
此时你容器中 /root/.cache/openpi/openpi-assets/checkpoints
目录下应该出现了一个 pi05_droid/
文件夹,里面存放的是模型权重:
[可选] Step5. 保存容器 & 打包镜像
使用下面的命令在 宿主机 上执行容器保存,附上作者信息和备注:
$ sudo docker commit -a "GaohaoZhou" -m "openpi" openpi openpi
$ sudo docker images
使用下面命令将镜像打包:
$ sudo docker save -o openpi.tar openpi
🌟 使用噪声进行推理
启动客户端 client
回到 /workspace/openpi
目录下执行下面的命令:
$ cd /workspace/openpi
$ python examples/simple_client/main.py --env DROID
启动服务端
你可以新开一个终端或者使用 tmux 创建一个终端来启动服务端,这里以新建一个终端的方式为例。
首先进入容器:
$ sudo docker exec -it openpi /bin/bash
然后进入/workspace/openpi
目录下执行下面的命令:
$ cd /workspace/openpi
$ python scripts/serve_policy.py --env DROID
运行后可以看到客户端和服务端的输出:
Terminal | |
---|---|
Server | ![]() |
Client | ![]() |
显存和硬件资源信息:
🤖 Aloha 仿真推理
【待 Jax 与 CUDA 冲突解决后完善】