ddyw

ddyw 收录面向机器人领域的开源模型与工具包。

目前，本仓库包含三类模型：

• 基于流匹配（flow matching）的视觉-语言-动作模型（VLA）。
• 基于快速动作分词器的自回归 VLA。
• 升级版基座 VLA，具有更强的开放世界泛化能力，采用改进的预训练策略训练。注意：在本仓库中，目前训练与推理仅支持流匹配输出头（flow matching head）。

对于所有模型，我们都提供了在 1 万小时以上机器人数据上预训练的基座模型检查点，并附带开箱即用或将其微调到自有数据集的示例。

这是一次试验：基座模型是为我们自己的机器人开发的，这些机器人与 ALOHA、DROID 等广泛使用的平台有所不同。我们乐观地认为，研究者和实践者能够发挥创意，将基座模型适配到自己的平台上，但我们并不指望每一次尝试都能成功。简而言之：基座模型可能适合你，也可能不适合，但欢迎你亲自尝试！

环境要求

要运行本仓库中的模型，你需要一块至少满足以下规格的 NVIDIA GPU。以下估算基于单块 GPU；你也可以通过在训练配置中设置 fsdp_devices 来使用多 GPU 模型并行，从而降低单卡显存需求。另请注意：当前训练脚本尚不支持多节点训练。

模式	所需显存	示例 GPU
推理	> 8 GB	RTX 4090
微调（LoRA）	> 22.5 GB	RTX 4090
微调（全量）	> 70 GB	A100 (80GB) / H100

本仓库已在 Ubuntu 22.04 上测试，目前不支持其他操作系统。

在你自己的数据上微调基座模型

我们以在 LIBERO 数据集上微调升级版基座 VLA 为例，介绍如何在自己的数据上微调基座模型。共分三步：

1. 将你的数据转换为 LeRobot 数据集（我们用它进行训练）
2. 定义训练配置并运行训练
3. 启动策略服务器并运行推理

1. 将你的数据转换为 LeRobot 数据集

我们在 examples/libero/convert_libero_data_to_lerobot.py 中提供了一个将 LIBERO 数据转换为 LeRobot 数据集的最小示例脚本。你可以轻松修改它来转换自己的数据！原始 LIBERO 数据集可从这里下载，然后运行脚本：

uv run examples/libero/convert_libero_data_to_lerobot.py --data_dir /path/to/your/libero/data

注意： 如果你只想在 LIBERO 上微调，可以跳过这一步，因为我们的 LIBERO 微调配置已指向一个预先转换好的 LIBERO 数据集。这一步只是作为一个示例，供你改编用于自己的数据。

2. 定义训练配置并启动训练

要在自己的数据上微调基座模型，你需要定义数据处理和训练的配置。我们为 LIBERO 提供了带详细注释的示例配置，你可以针对自己的数据集进行修改：

• LiberoInputs 和 LiberoOutputs：定义 LIBERO 环境与模型之间的数据映射，训练和推理都会用到。
• LeRobotLiberoDataConfig：定义如何处理来自 LeRobot 数据集的原始 LIBERO 数据用于训练。
• TrainConfig：定义微调超参数、数据配置和权重加载器。

我们提供了流匹配 VLA、自回归 VLA 和升级版基座 VLA 在 LIBERO 数据上的微调示例配置（见训练配置文件）。

在开始训练之前，需要先计算训练数据的归一化统计量。使用你的训练配置名称运行以下脚本：

uv run scripts/compute_norm_stats.py --config-name pi05_libero

现在可以用以下命令启动训练（如果你用同一配置重新微调，--overwrite 标志用于覆盖已有检查点）：

XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run scripts/train.py pi05_libero --exp-name=my_experiment --overwrite

该命令会将训练进度输出到控制台，并把检查点保存到 checkpoints 目录。你还可以在 Weights & Biases 面板上监控训练进度。为了最大限度利用 GPU 显存，请在运行训练前设置 XLA_PYTHON_CLIENT_MEM_FRACTION=0.9——这允许 JAX 使用最多 90% 的 GPU 显存（默认为 75%）。

注意： 我们提供从预训练中重载状态/动作归一化统计量的功能。当你在一个属于我们预训练数据分布的机器人上微调新任务时，这可能带来收益。有关如何重载归一化统计量的更多细节，请参阅 norm_stats.md。

3. 启动策略服务器并运行推理

训练完成后，我们可以启动一个策略服务器，然后在 LIBERO 评测脚本中查询它来运行推理。启动模型服务器很简单（本例使用第 20000 次迭代的检查点，可按需修改）：

uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi05_libero --policy.dir=checkpoints/pi05_libero/my_experiment/20000

这会启动一个监听 8000 端口的服务器，等待观测数据发送过来。随后我们可以运行评测脚本（或机器人运行时）来查询该服务器。

对于运行 LIBERO 评测，我们提供（并推荐使用）一个将策略服务器和评测脚本整合在一起的 Docker 化工作流。详见 LIBERO README。

如果你想在自有机器人运行时中嵌入策略服务器调用，我们在远程推理文档中提供了一个最小示例。

更多示例

我们在以下 README 中提供了在 ALOHA 平台上微调和运行推理的更多示例：

• ALOHA 仿真
• ALOHA 实机
• UR5

PyTorch 支持

ddyw 现在在原有 JAX 版本之外，提供了流匹配 VLA 和升级版基座 VLA 的 PyTorch 实现！PyTorch 实现已在 LIBERO 基准上完成验证（推理和微调）。以下功能目前尚不支持（未来可能会变化）：

• 自回归 VLA 模型
• 混合精度训练
• FSDP（全分片数据并行）训练
• LoRA（低秩适应）训练
• 训练期间的 EMA（指数移动平均）权重

设置

1. 确保所有依赖都是最新版本：uv sync

2. 再次确认已安装 transformers 4.53.2：uv pip show transformers

3. 应用 transformers 库补丁：

   cp -r ./src/openpi/models_pytorch/transformers_replace/* .venv/lib/python3.11/site-packages/transformers/

此操作会用必要的模型改动覆盖 transformers 库中的若干文件：1) 支持 AdaRMS，2) 正确控制激活值的精度，3) 允许在不更新的情况下使用 KV cache。

警告：使用 uv 默认的链接模式（硬链接）时，这会永久影响 uv 缓存中的 transformers 库，意味着更改在重新安装 transformers 后依然存在，甚至可能传播到其他使用 transformers 的项目。要完全撤销此操作，必须运行 uv cache clean transformers。

将 JAX 模型转换为 PyTorch

要将 JAX 模型检查点转换为 PyTorch 格式：

uv run examples/convert_jax_model_to_pytorch.py \
    --checkpoint_dir /path/to/jax/checkpoint \
    --config_name <config name> \
    --output_path /path/to/converted/pytorch/checkpoint

使用 PyTorch 运行推理

PyTorch 实现使用与 JAX 版本相同的 API——只需将检查点路径改为指向转换后的 PyTorch 模型：

from openpi.training import config as _config
from openpi.policies import policy_config
from openpi.shared import download

config = _config.get_config("pi05_droid")
checkpoint_dir = "/path/to/converted/pytorch/checkpoint"

# 创建训练好的策略（自动检测 PyTorch 格式）
policy = policy_config.create_trained_policy(config, checkpoint_dir)

# 运行推理（与 JAX 相同的 API）
action_chunk = policy.infer(example)["actions"]

PyTorch 策略服务器

策略服务器对 PyTorch 模型的工作方式完全相同——只需指向转换后的检查点目录：

uv run scripts/serve_policy.py policy:checkpoint \
    --policy.config=pi05_droid \
    --policy.dir=/path/to/converted/pytorch/checkpoint

使用 PyTorch 微调

使用 PyTorch 微调模型：

1. 将 JAX 基座模型转换为 PyTorch 格式：

   uv run examples/convert_jax_model_to_pytorch.py \
       --config_name <config name> \
       --checkpoint_dir /path/to/jax/base/model \
       --output_path /path/to/pytorch/base/model

2. 在你的配置中使用 pytorch_weight_path 指定转换后的 PyTorch 模型路径

3. 使用以下模式之一启动训练：

# 单卡训练：
uv run scripts/train_pytorch.py <config_name> --exp_name <run_name> --save_interval <interval>

# 示例：
uv run scripts/train_pytorch.py debug --exp_name pytorch_test
uv run scripts/train_pytorch.py debug --exp_name pytorch_test --resume  # 从最新检查点恢复

# 多卡训练（单节点）：
uv run torchrun --standalone --nnodes=1 --nproc_per_node=<num_gpus> scripts/train_pytorch.py <config_name> --exp_name <run_name>

# 示例：
uv run torchrun --standalone --nnodes=1 --nproc_per_node=2 scripts/train_pytorch.py pi0_aloha_sim --exp_name pytorch_ddp_test
uv run torchrun --standalone --nnodes=1 --nproc_per_node=2 scripts/train_pytorch.py pi0_aloha_sim --exp_name pytorch_ddp_test --resume

# 多节点训练：
uv run torchrun \
    --nnodes=<num_nodes> \
    --nproc_per_node=<gpus_per_node> \
    --node_rank=<rank_of_node> \
    --master_addr=<master_ip> \
    --master_port=<port> \
    scripts/train_pytorch.py <config_name> --exp_name=<run_name> --save_interval <interval>

精度设置

JAX 与 PyTorch 实现的精度处理方式如下：

JAX：

1. 推理：大多数权重和计算使用 bfloat16，少量计算为稳定性使用 float32
2. 训练：默认混合精度：权重和梯度为 float32，（大部分）激活值和计算为 bfloat16。可以在配置中将 dtype 设为 float32 来使用全量 float32 训练。

PyTorch：

1. 推理：与 JAX 一致——大多数权重和计算使用 bfloat16，少量权重为稳定性转换为 float32
2. 训练：支持全量 bfloat16（默认）或全量 float32。可以通过在配置中设置 pytorch_training_precision 来切换。bfloat16 显存占用更少，但相比 float32 损失更高。混合精度暂不支持。

使用 torch.compile 时，JAX 与 PyTorch 的推理速度相当。

常见问题

我们会在这里收集常见问题及其解决方案。遇到问题时请先查看这里。如果找不到解决方案，请在仓库中提交 issue（参见此处的指南）。

问题	解决方案
uv sync 因依赖冲突失败	尝试删除虚拟环境目录（rm -rf .venv）后重新运行 uv sync。如果问题依旧，检查 uv 是否为最新版本（uv self update）。
训练时 GPU 显存不足	运行训练前请确保已设置 XLA_PYTHON_CLIENT_MEM_FRACTION=0.9（或更高），以允许 JAX 使用更多 GPU 显存。也可以使用 --fsdp-devices <n>（<n> 为你的 GPU 数量）启用全分片数据并行，以训练变慢为代价降低显存占用（具体变慢程度取决于你的环境）。如果显存仍然不足，可以考虑关闭 EMA。
策略服务器连接错误	检查服务器是否正在运行并监听预期端口。验证客户端与服务器之间的网络连通性和防火墙设置。
训练时缺少归一化统计量错误	在开始训练前，用你的配置名称运行 scripts/compute_norm_stats.py。
数据集下载失败	检查网络连接。对于 HuggingFace 数据集，请确保已登录（huggingface-cli login）。
CUDA/GPU 错误	确认 NVIDIA 驱动已正确安装。对于 Docker，确保已安装 nvidia-container-toolkit。检查 GPU 兼容性。系统层面无需安装 CUDA 库——它们会通过 uv 安装。如果遇到 CUDA 问题，甚至可以尝试卸载系统 CUDA 库，因为系统库有时会引起冲突。
运行示例时出现导入错误	确保已通过 uv sync 安装所有依赖。部分示例可能在各自 README 中列出了额外依赖。
动作维度不匹配	检查你的数据处理变换是否与你机器人的预期输入/输出维度匹配。查看策略类中的动作空间定义。
训练损失发散	检查数据集 norm_stats.json 中的 q01、q99 和 std 值。某些很少使用的维度可能得到非常小的 q01、q99 或 std 值，导致归一化后出现巨大的状态和动作。可以手动调整归一化统计量作为临时解决办法。