LeRobot documentation

Praxis-1

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.7.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Praxis-1

praxis 是 Praxis-1 的 LeRobot 适配器,Praxis-1 是 Runway 推出的面向机器人的开放权重世界动作模型。Praxis-1 是一个视频-动作模型:其块因果视频扩散主干(一个 27B 参数的混合专家 Transformer,每个 token 约激活 3B 参数)对两个机架相机的未来 1.6 秒视频进行去噪,同时动作专家读取主干的隐藏状态,对匹配的 24 个笛卡尔命令在 15 Hz 频率下进行去噪。Wan2.2 视频自编码器和 Qwen3.5-2B 文本编码器保持冻结,并来自其公开权重。发布的检查点 praxis-1-droid 驱动单臂 DROID 机架:一个配备外部相机和腕部相机的 Franka 机械臂。一个更早的、未发布的 DROID 检查点具有相同的架构但使用不同的权重和另一个相机画布;检查点的 config.json 可以区分它们(参见“检查点”)。praxis-1-yam 是用于两个 YAM 机械臂的双臂检查点,采用相同的布局(参见“双臂 YAM 检查点”)。praxis-1-base 是视频主干和动作专家主干,不包含特定机架的头部,是微调新单臂机架的起点(参见“微调基础检查点”)。

资源:

资源链接用于
公告runway.com/news/praxis-1-is-now-available方法概述
权重Praxis-1 数据集runwayml/praxis-1-droid-lerobot 以及其他转换后的检查点
代码runwayml/praxis-1-oss原始推理、模拟和训练
模拟器PyPI 上的 praxis-simMuJoCo 测试装置及其评估套件
Wan2.2 VAEWan-AI/Wan2.2-TI2V-5BWan2.2_VAE.pth(冻结)
文本编码器Qwen/Qwen3.5-2B冻结

模型代码位于仓库内的 lerobot/policies/praxis/praxis_core(从 praxis 仓库的 8dad1a4 提交版本引入,并添加了 praxis-1-droid 的 canvas 规则,重写了导入路径;文件头部注明了上游提交)。MuJoCo 模拟是 PyPI 上的 praxis-sim 包(在 praxis 仓库的 packages/praxis-sim 中开发,不包含 torch);LeRobot 没有为其提供环境。

安装

pip install -e ".[praxis,dataset]"   # the policy; `lerobot-eval` and `lerobot-train` also need `dataset`

转换后的检查点已托管在 Hub 上:--policy.path=runwayml/praxis-1-droid-lerobot 用于加载已发布的单臂检查点(54 GB,只需下载一次)。runwayml/praxis-1-yam-lerobot 是双臂检查点。runwayml/praxis-1-base-lerobot 是基础模型,需要先运行 init_heads.py。

该策略运行在 LeRobot 0.7 锁定的版本上(撰写本文时为 0.7.1:torch 2.15.0 搭配 CUDA 13.2,transformers 5.22.2,triton 3.9)。以下结果是在该技术栈(使用 transformers 5.22.1 和 5.22.2,两者生成的命令相同)以及 LeRobot 0.6 的技术栈(torch 2.11.0 搭配 CUDA 12.8,transformers 5.5.4)上测量的;在每种情况下,这些命令都与 praxis 仓库中的命令按位相等。

检查点验证配置中的可选内核。flash-attn-4 提供检查点配置中的 flash 注意力后端(使用 FlashAttention-4 CuTe 内核的 flex attention);若未安装,策略将使用 flex attention 的 Triton 内核。flash-linear-attention 和 causal-conv1d 是 Qwen3.5 线性注意力层的高速内核;若未安装,transformers 将运行其 PyTorch 实现。这三者均会改变输出的最后几位,而非模型本身:

pip install "flash-attn-4[cu13]==4.0.0b36"        # the CUDA 13 build for torch 2.15; "flash-attn-4==4.0.0b36" on a CUDA 12 torch
pip install "flash-linear-attention==0.6.0"       # 0.5.0 does not run on triton 3.9
CAUSAL_CONV1D_FORCE_BUILD=TRUE MAX_JOBS=16 pip install --no-build-isolation "causal-conv1d==1.6.1"   # compiles against torch; the praxis repository's setup

推理需要一块具有 80 GB 内存的 GPU(权重在 bfloat16 下为 54 GB)。

在数据中心 Blackwell GPU 上(B200 和 GB200 为 SM100,B300 和 GB300 为 SM103;Jetson Thor 为 SM110,具有相同的 FlashAttention-4 前向传播):

  • B300 and GB300 need a torch build for CUDA 12.9 or later, such as the torch 2.15.0+cu132 that LeRobot 0.7 locks. The MoE experts call torch._grouped_mm, and the build of torch 2.11.0+cu128 (LeRobot 0.6) has no SM103 kernel for it: the first policy call fails with RuntimeError: cutlass cannot run, error 7, as the praxis repository’s pipeline does. The B200 runs on both stacks.
  • FlashAttention-4 在此处仅接受以 256 个 token 为块大小的块稀疏掩码(否则会抛出 ValueError: Block sparse tensorsforward have block size 128, which must be a multiple of 256)。 使用 flash 后端时,策略会在这些 GPU 上以 256 个 token 为块大小构建主干掩码,而不是检查点中 backbone.attention.flex_block_size 指定的 128。块大小仅用于对 token 对进行分组:掩码函数仍然决定每一对 token 的掩码值。praxis 仓库在检查点的 config.json 中设置 "flex_block_size": 256 来运行相同的配置。triton 后端则保留检查点中的 128。

在 RTX Blackwell GPU(计算能力 12.x;RTX PRO 6000 Blackwell 拥有权重所需的 96 GB 显存)上,FlashAttention-4 在任何块大小下均不支持块稀疏前向传播(使用检查点的 128 个 token 块时,会抛出 Block sparsity requires sparse_block_size[1]=128 to match tile_n=64),因此 flash 后端无法在这些 GPU 上运行。attention_backend="auto" 会在这类 GPU 上选择 triton,即使已安装 flash-attn-4 也是如此,而 "flash" 在构建策略时会引发 ValueError。praxis 仓库在首次调用时若使用检查点的 flash 会失败,而在检查点的 config.json 中设置 "backend": "triton" 后则可正常运行。这两个后端在最低有效位上存在差异,因此请将这些 GPU 上的运行结果与在其他位置使用 triton 的运行结果进行对比。在 RTX PRO 6000 上其他方面均无变化:下方两个技术栈的 torch 构建版本均包含 SM120 内核(包括 MoE 的 torch._grouped_mm),且使用 praxis-1-droid 的设置进行推理时,峰值显存占用约为 62 GB。

这些规则遵循 FlashAttention-4 自身的检查。以下奇偶校验结果是在 GB300、B200 和 RTX PRO 6000 Blackwell 上测量的。

transformers 5.4.0 至 5.5.4(LeRobot 0.6;已在 5.6.0 中修复,且 LeRobot 0.7 锁定 5.22)在加载 Hub 检查点 Qwen/Qwen3.5-2B 时,会将全部 320 个语言模型张量报告为缺失(它在匹配前重命名了存储的 model.language_model.* 名称),并静默地以随机值初始化它们。策略的文本编码器会检测到这一情况,并按存储的名称加载张量,同时在日志中发出警告。若文本编码器保持随机初始化,策略虽然仍能移动,但极少能成功(在我们首次运行中,56 个 guard 回合里仅成功 7 次)。

输入和输出

功能形状含义
observation.images.exteriorRGB固定摄像头。任意尺寸;预处理器会根据检查点生成对应的图像块(对于praxis-1-droid,采用抗锯齿适配规则,尺寸为 224 x 416;对于较早的 DROID 检查点,尺寸为 256 x 416)。键由exterior_image_key设置。
observation.images.wristRGB手腕摄像头。由 wrist_image_key 设置的关键字。
observation.state(7,)[x, y, z, roll, pitch, yaw, gripper]:机械臂基坐标系下的 Franka 法兰位姿(单位为米,采用 DROID 约定的外旋 x-y-z 欧拉角,单位为弧度)以及夹爪开合度(取值范围为 [0, 1])。
taskstr指令,原样提供给文本编码器。
action(7,)相同的七个值:下一个 15 Hz 时间步的绝对目标位姿和夹爪状态。
observation.scene_seed(可选)(1,)该集的场景种子。块 k 的采样器种子为 seed + seed_stride * scene_seed + k。

策略会维护自己的历史:一个环形缓冲区,存储批次中每个环境最近 25 个观测值(即已执行的块加上当前帧),并将其作为上下文提供给下一次调用,因此 n_obs_steps 保持为 1,并且必须在每个 15 Hz 的 tick 上调用 select_action,传入该 tick 的观测值。上下文窗口的第一次调用返回 21 个命令,之后的每次调用返回 24 个。如果 n_action_steps 小于 24,则每次调用都是没有上下文的冷启动,与 praxis 中的情况相同。drop_queued_actions(例如在 rollout 中途出现新指令)也会重启时间上下文,因为模型以整个前一个块已被执行为条件。执行完整块的调用方(predict_action_chunk)会将其上一次调用以来已执行的帧作为堆栈传递,即 (B, T, 3, H, W) 的相机图像和 (B, T, 7) 的状态,最旧的在前;那么下一个块就是逐 tick 调用 rollout 的那个块。带有一帧的块调用会重新开始上下文。

处理器

预处理器以与 praxis 数据和 serve 代码相同的方式准备输入,并作为已保存的、可检查的步骤:

步骤是否
PraxisTileViewsProcessorStep每个相机视角都会成为画布上的一块图块,其处理规则由检查点(video.antialias)决定。praxis-1-droid:中心裁剪至原生 720 x 1280 宽高比,进行一次抗锯齿双线性缩放至 234 x 416(如果视角已以该尺寸发送,则不进行缩放),中心裁剪至 224 x 416。较早的 DROID 检查点:中心裁剪至图块宽高比,进行一次普通双线性缩放至 256 x 416。使用 uint8 色阶;如果视角已处于图块尺寸,则直接通过。
PraxisPoseProcessorStepDROID 位姿 [x, y, z, roll, pitch, yaw, gripper] 变为 praxis 位姿8 [x, y, z, qw, qx, qy, qz, gripper](状态,以及训练批次中的动作)。
PraxisTokenizerProcessorStep指令会被转换为 Qwen3.5-2B 文本编码器的 token ids 和 attention mask(observation.language.tokens、.attention_mask),其分词方式与 praxis 文本编码器的分词方式一致。冻结的编码器及其嵌入缓存保留在 policy 中。
PraxisCommandProcessorStep(postprocessor) 策略的命令、[x, y, z, rotation 3 x 3] 形式的 praxis 线位姿以及夹爪,在 CPU 上转换为 DROID 的 7 个值。

在它们之间,策略(policy)接收 tiles 和 pose8 行,并输出 wire poses,这正是 Session 所消费和产生的内容。q99 动作归一化、均值/标准差本体感觉归一化以及相对动作解码都保留在会话内部:它们锚定在每个执行块开始时的 pose,会话会在 rollout 过程中对其进行跟踪,而它们的统计量来自检查点(PraxisConfig.model_config["normalization"]);默认的归一化器步骤是恒等变换。策略还接受原始的 7 值 pose、原始相机帧以及 task 字符串,因此可以在没有 pipeline 的情况下调用它:tile 尺寸的帧被视为 tiles,其他尺寸的帧则通过检查点的 canvas 规则处理(praxis-1-droid 的 fit 规则不是幂等的,因此一个 tile 永远不会经过它两次),而没有 token 的 task 会由策略进行分词。tiles 会保留在 pipeline 将其放置的设备上,直到 autoencoder 读取它们,而命令则返回到 CPU:没有任何内容会两次跨越设备边界。

检查点

检查点Tiles调整大小推荐采样器(praxis-1-droid 的 PraxisConfig 默认值)
praxis-1-droid(已发布)224 x 416抗锯齿适配规则(video.antialias: true)2 步 Euler 采样,flow shift 12.5,视频上的 text guidance 为 4.0,动作上无 guidance(1.0),1 个 context chunk
YAM,praxis-1-yam(双手操作,已发布)256 x 416,三个图块YAM 拟合规则 (video.antialias: true)2 步,flow shift 1.0,无文本引导,1 个上下文块(由转换器写入)
基础模型,praxis-1-base(无头部)由 init_heads.py 设置由 init_heads.py 设置微调的选择;DROID 装置采用 praxis-1-droid 的
之前的 DROID 检查点(未发布)256 x 416plain bilinear resize2 步,shift 8.0,无文本引导,3 个上下文块(praxis);以下结果使用了 shift 1.0、动作引导 3.0 和 1 个上下文块

画布来自检查点的 config.json(video.tile_height、video.antialias),通过转换器传递;采样器设置是 PraxisConfig 的字段,因此运行较早的 DROID 检查点时会传入其自身的(--policy.t_shift=8.0 --policy.cfg_video_scale=1.0 --policy.max_context_chunks=3)。

转换检查点

Praxis 检查点是一个包含 config.json、17 个 safetensors 分片、model.safetensors.index.json 和 SHA256SUMS 的目录。将其转换为 LeRobot 策略目录:

python examples/praxis/convert_checkpoint.py \
    --checkpoint /path/to/praxis-1-droid --out outputs/praxis-1-droid --link

该装置源自检查点的 embodiment:双手操作检查点会获取其三个相机键(可通过 --camera-keys 重命名)以及其检查点的服务设置。基础检查点会获得一个不包含处理器的目录,由 init_heads.py 补全。runwayml/praxis-1-droid-lerobot 及其他已转换的检查点均为该脚本的输出。

--link 通过符号链接而非复制来管理 54 GB 的分片。该目录包含 LeRobot 的 config.json、分片、索引及其自身的 SHA256SUMS(在其中运行 sha256sum -c 可通过校验)(LeRobot 策略以 transformers 布局读写分片 safetensors;PraxisPolicy.save_pretrained 会写入 5 GB 的分片),以及两个处理器文件。--vae-id 和 --text-encoder-id 用于设置冻结部分的来源(默认为 Hub id)。

结果

以下闭环数据来自 praxis 仓库中的 MuJoCo 装置(praxis-sim 软件包:Franka + Robotiq 2F-85,512 x 288 摄像头,15 Hz)及其固定的评估套件;LeRobot 中某一集(episode)的种子即为其场景种子。通过 lerobot-eval 运行这些套件的环境不属于 LeRobot 的一部分。

在采样过程中,策略会固定 praxis 运行所采用的数值设置(禁用 cuDNN benchmark 模式、MoE 路由器使用 float32 矩阵乘法、允许 cuDNN 使用 TF32),并在之后恢复主机设置,因此 --tf32、torch.backends.cudnn.benchmark 和批处理大小不会改变其输出。

droid_guard 测试套件(7 个 DROID 任务 x 8 个场景,56 个回合,采样器种子 0),使用早期未发布的 DROID 检查点(256 x 416 图块),采样器设置采用当时 praxis 仓库的配置(2 步 Euler,flow shift 1.0,action guidance 3.0,无 video guidance,1 个 context chunk),并使用其 flash attention 后端,在单张 H100 80 GB 上运行(约 10 分钟)。相同的回合以五种方式运行:praxis 仓库自带的 runner 在其自身环境中(torch 2.13,transformers 5.2.0,causal-conv1d);相同的 runner 和 lerobot-eval 在 LeRobot 0.6 的栈上(torch 2.11,transformers 5.5.4);以及两者再次在 LeRobot 0.7 锁定的栈上运行(torch 2.15 搭配 CUDA 13.2,triton 3.9,flash-linear-attention 0.6.0;当时的 transformers 5.16.1)。使用 transformers 5.22.1 重复 torch 2.15 的运行,对于 praxis runner 和 lerobot-eval 均产生逐位相同的命令以及相同的 46/56 回合结果;而 5.22.2(LeRobot 0.7.1 锁定的版本)为 praxis-1-droid 和 YAM 检查点产生相同的命令:transformers 版本不影响数值。

任务praxis runner, praxis envpraxis runner, torch 2.11lerobot-eval,torch 2.11praxis runner, torch 2.15lerobot-eval,torch 2.15
lift_cube7/87/86/87/88/8
cube_on_plate8/88/88/88/88/8
block_in_bowl8/88/88/88/88/8
block_out_of_bowl7/86/87/87/86/8
move_block7/87/86/87/87/8
stack_block7/88/87/88/86/8
press_button1/80/83/81/83/8
总计45/56 80 %44/56 79 %45/56 80 %46/56 82 %46/56 82 %

在 56 个回合中配对(praxis_sim.compare),lerobot-eval 减去参考值:

对比差异95% 自举区间胜 / 负McNemar p
torch 2.11: lerobot-eval 与 praxis runner、praxis env+0.0 %[-10.7 %, +10.7 %]5 / 51.00
torch 2.11:lerobot-eval 与 praxis runner,相同的软件栈+1.8 %[-8.9 %, +12.5 %]6 / 51.00
torch 2.15:lerobot-eval 与 praxis runner、praxis env+1.8 %[-10.7 %, +14.3 %]6 / 51.00
torch 2.15:lerobot-eval 与 praxis runner,相同技术栈+0.0 %[-10.7 %, +10.7 %]4 / 41.00

每个区间都包含零:LeRobot 策略和环境复现了 praxis 仓库的结果,且软件栈(内核、transformers 版本)或测试框架的变化均未超出逐回合噪声范围地改变结果。praxis 运行器本身在不同软件栈之间的移动幅度也相同(56 中的 44 到 46)。press_button 对于该检查点在所有配置下都很难。

在输入相同的情况下,该策略在两个栈上均与 praxis 流水线按位相等(详见下文);而在两个栈之间,相同输入的命令差异最大为 0.04(flash)和 0.02(triton),正是这些差异导致了单个 episode 的变化。

praxis-1-droid,已发布的检查点

使用 praxis-1-droid 的相同 56 个剧集,在其推荐设置下(224 x 416 抗锯齿瓦片,2 步 Euler,flow shift 12.5,视频引导 4.0,无动作引导,1 个上下文块)以及 flash 后端,在单张 H100 80 GB 上运行:praxis 仓库的 runner(其 main 配合当前采样器,在 praxis 环境中:torch 2.13,transformers 5.2.0)、基于 LeRobot 0.6 技术栈的 lerobot-eval(torch 2.11,transformers 5.5.4),以及基于 LeRobot 0.7 锁定技术栈的 runner 和 lerobot-eval(torch 2.15 配合 CUDA 13.2,transformers 5.22.1)。

任务praxis runner, praxis envlerobot-eval,torch 2.11praxis runner, torch 2.15lerobot-eval,torch 2.15
lift_cube7/88/88/88/8
cube_on_plate8/87/88/88/8
block_in_bowl8/88/88/88/8
block_out_of_bowl8/88/87/88/8
move_block7/88/88/88/8
stack_block8/87/88/86/8
press_button8/88/87/88/8
总计54/56 96 %54/56 96 %54/56 96 %54/56 96 %

在 56 个回合的配对比较中,lerobot-eval 减去 praxis runner 在两个技术栈上的得分均为 +0.0 分(95% 自助法区间 [-7.1%, +7.1%],各赢 2 局 / 输 2 局,McNemar p 值均为 1.00)。在两个技术栈上发生胜负翻转的回合并不相同:torch 2.15 的内核使指令最多偏移 0.02,少数处于临界状态的回合随之改变结果。praxis-1-droid 使用其配置能够解决 press_button,而使用其配置的早期 DROID 检查点大多未能解决该任务。

与 praxis 仓库保持一致

  • 在单进程中对相同观测值,LeRobot 策略与 praxis Session 生成的文本嵌入、VAE 潜变量、本体感知、采样动作和指令在比特级别上完全相等(加载的权重相同)。tests/policies/praxis/test_praxis.py 在 CPU 上使用一个小型随机模型重复执行此检查。
  • 在跨进程场景中,通过原始 praxis 包运行的 praxis 流水线与 LeRobot 策略在真实检查点上生成的命令在位级上完全一致,无论是在 flash 还是 triton 后端(这两个后端之间的命令差异最大为 0.055),无论是在 LeRobot 0.6 的依赖栈(torch 2.11.0+cu128, transformers 5.5.4, flash-attn-4 4.0.0b36)上,还是在 LeRobot 0.7 锁定的依赖栈(torch 2.15.0+cu132, transformers 5.22.1 和 5.22.2, triton 3.9, flash-attn-4 4.0.0b36)上。参考实现本身在两个依赖栈之间也存在差异(在 flash 上最大为 0.019,在 triton 上为 0.015);策略在每一个栈上都逐位跟随参考实现。
  • praxis-1-droid:针对 praxis 仓库 main 分支(含其当前采样器)的 praxis 流水线进行的跨进程检查,在两种后端和两种技术栈下,以八种方式调用策略时均实现逐位相等:通过其保存的处理器流水线(如 lerobot-eval 所做的那样),在原始帧和 7 值位姿上运行;不使用流水线,在一个批次中处理两个环境(不同的指令和场景种子);通过 predict_action_chunk(冷启动,以及给定已执行帧作为堆栈的下一个块);设置 n_action_steps=12(每 12 个命令一个块,以已执行帧为条件);开启动作引导(cfg_scale=3);以及跨越 70 个命令,包含三个上下文块(深度 0 到 3)。该检查在 lerobot-eval 的默认 TF32 和 cuDNN benchmark 标志下依然成立:策略在采样期间固定 praxis 的数值,并在完成后将标志交还给主机。无论帧所在的设备为何,tile 步骤都在 CPU 上计算画布:CUDA 抗锯齿缩放会在约三分之一的像素上产生不同级别的舍入,这导致在 tile 步骤执行之前,命令最多偏移 0.02。
  • 训练:在发布的权重上,praxis 仓库中的 PraxisPolicy.forward 和 praxis.train.objective 对于相同的窗口和随机抽样,逐位地给出相同的损失(以及视频和动作项),并且在动作专家的所有 200 个张量上给出相同的梯度。报告的 MoE 平衡指标在任意一侧的运行之间相差一或两个 float32 ulp(平衡统计量中的非确定性归约;它以 1e-4 的权重进入损失)。
  • GB300(SM103,aarch64 主机):在 LeRobot 0.7 锁定的软件栈(torch 2.15.0+cu132、cuDNN 9.26、transformers 5.22.1、triton 3.9、flash-attn-4 4.0.0b36、flash-linear-attention 0.6.0)上,针对 praxis 仓库的 main 分支进行的跨进程检查在两个后端(flash,两侧均使用 256 个 token 的掩码块)上逐位相等:对于 praxis-1-droid,在七种策略调用方式下(lerobot-eval 标志下的已保存流水线、原始 720 x 1280 帧、一个批次中的两个环境、带帧堆栈的 predict_action_chunk、n_action_steps=12、cfg_scale=3,以及带有三个上下文块的 94 条指令);对于 praxis-1-yam,通过流水线、在原始帧上并带有两个环境;以及对于从 praxis-1-base 重建的 DROID 装置。forward 在发布的权重上,逐位给出 praxis 目标的损失以及所有 200 个动作专家张量的梯度。
  • RTX PRO 6000 Blackwell(SM120,96 GB,x86_64 主机),两侧均使用 triton:一侧是 LeRobot 0.6 的依赖栈(torch 2.11.0+cu128,cuDNN 9.19,transformers 5.5.4,triton 3.6,flash-linear-attention 0.5.0),另一侧是 LeRobot 0.7 锁定的依赖栈(torch 2.15.0+cu132,cuDNN 9.26,transformers 5.22.1,triton 3.9,flash-linear-attention 0.6.0 以及 causal-conv1d 1.6.1,因此 Qwen3.5 运行其内核),两者均安装了 flash-attn-4 4.0.0b36。针对 praxis 仓库的 main 进行的跨进程检查,在 GB300 所体现的方面上是逐位相等的:praxis-1-droid 在七个(463 条命令,上下文深度 0 到 3)中,praxis-1-yam 在三个中,以及从 praxis-1-base 重建的 DROID 装置(其六个头部张量及其配置与 praxis 仓库的 init_heads 相等)。forward 在两个依赖栈上逐位给出 praxis 目标的损失以及所有 200 个动作专家张量的梯度,两侧均冻结骨干网络(其 54 GB 的梯度无法与权重同时容纳)。在 LeRobot 0.6 的依赖栈上,praxis 侧也需要上述警告的 Qwen3.5 修复:praxis-1 固定了 transformers 5.2.0,而在 5.5.4 下其文本编码器保持随机状态。在 LeRobot 0.6 的依赖栈上使用 praxis-1-droid 对 droid_guard 进行闭环测试:lerobot-eval 为 53/56(94.6 %),praxis 运行器为 54/56(96.4 %,与 H100 和 GB300 一致);配对比较,-1.8 分,95 % 区间 [-8.9, +5.4],胜 2 / 负 3,McNemar p 1.00。通过 lerobot-train 进行的 LoRA 微调(r 16,alpha 32,默认目标,每步一个窗口)在单张 GPU 上运行:5.5M 可训练参数,每步 1.1 秒,62 GB。
  • B200(SM100,x86_64 主机):在 LeRobot 0.6 的依赖栈(torch 2.11.0+cu128,cuDNN 9.19,transformers 5.5.4,triton 3.6)以及 LeRobot 0.7 锁定的依赖栈(torch 2.15.0+cu132,cuDNN 9.26,transformers 5.22.1,triton 3.9,flash-linear-attention 0.6.0)上,两者均使用 flash-attn-4 4.0.0b36,针对 praxis 仓库的 main 进行的跨进程检查在两个后端(flash,两侧均使用 256 个 token 的掩码块)上均实现逐位相等,情况与 GB300 一致:praxis-1-droid 在七项测试中(每个后端和依赖栈 463 条命令,上下文深度 0 到 3),praxis-1-yam 在三项测试中,以及从 praxis-1-base 重建的 DROID 装置(其六个头部张量和配置与 praxis 仓库的 init_heads 相等)。两个后端在命令上的差异最大为 0.19。forward 在两个后端和两个依赖栈上,针对已训练的主干网络,逐位给出 praxis 目标的损失、其视频和动作项以及所有 200 个动作专家张量的梯度。通过 lerobot-train 进行的 LoRA 微调可在单张 GPU 上运行(每步 1.2 秒,62 GB),并且该适配器及其合并的检查点可在 lerobot-eval 中加载。
  • B200 上的闭环测试(droid_guard、praxis-1-droid、flash):praxis runner 在 torch 2.15 栈上得分为 54/56,在 torch 2.11 上为 55/56;lerobot-eval 分别为 53/56 和 51/56(配对比较:-1.8 分,95% 置信区间 [-8.9, +5.4];以及 -7.1 分,[-14.3, -1.8],胜 0 / 负 4,McNemar p 0.125)。每个测试框架逐集重复执行,因此差距并非运行间的噪声,也不是策略本身的问题:lerobot-eval 以 float32 格式承载 DROID 的 7 个状态值和指令值(PraxisSimGymEnv,位姿和指令处理器步骤),而 praxis runner 则原样传递会话的位姿和指令。往返转换会使位姿产生约 1e-7 的偏移,在一个 episode 内,这种差异会累积成另一条轨迹。当 praxis runner 的状态和指令也经过相同的 float32 转换时,在两个栈上的全部 56 个 episode 中均得到与 lerobot-eval 一致的结果。
  • 针对 praxis 仓库自身的运行环境(torch 2.13.0+cu130、transformers 5.2.0、causal-conv1d),输出结果并非按位相等:内核存在差异,Qwen3.5 的嵌入在 bfloat16 下最大差异为 0.125,且动作的末位比特也随之变化。因此,仓库间的校验即为上述成对的闭环评估。
  • 该模拟器以 4 倍多重采样进行渲染,并且同一场景的多次渲染之间每帧会有少量像素差异,因此即使权重完全相同,同一策略的两个闭环 episode 在第一个 chunk 之后也会产生分歧。配对的 episode 通过统计方法进行比较,而非逐帧对比。

双手 YAM 检查点

praxis-1-yam("embodiment": "bimanual",runwayml/praxis-1-yam-lerobot)驱动两条 YAM 机械臂,配备一个顶部相机和两个腕部相机(16:10 ZED 视图,1920 x 1200)。运行的是同一个模型;检查点在专家头宽度(每条臂 20 个动作值:[xyz, rot6d, gripper])、三格画布以及归一化方面有所不同。LeRobot 适配器遵循 praxis 仓库中的 praxis/bimanual.py(以 praxis_core/bimanual.py 的形式内置):

密钥形状含义
observation.images.top_left, observation.images.wrist_left, observation.images.wrist_rightRGB画布顺序中的三个视图(camera_keys)。任意 16:10 尺寸;处理器使用 YAM 适配规则生成 256 x 416 的图块(视图共同适配一个 416 px 的框,然后进行 32 px 的裁剪)。
observation.state(14,)左臂的[x, y, z, roll, pitch, yaw, gripper],然后是右臂的,均在机械臂的世界坐标系中(pose_frame: rig,无法兰旋转)。16 个值(每臂 pose8)也会通过。
action(14,)相同的 14 个值:每个机械臂在下一个 15 Hz 节拍时的绝对目标。select_action 输出 26 个值(每个机械臂的末端位姿和夹爪状态);命令步骤会将它们进行转换。

The checkpoint’s serve settings (the praxis repository’s configs/serve_yam.yaml: 2 Euler steps, flow shift 1.0, no text guidance, 1 context chunk, the bimanual rig has no control mode) are what the converter writes into the config, so --policy.path=runwayml/praxis-1-yam-lerobot runs them. Prompts go to the text encoder verbatim; the checkpoint was trained on captions such as “Pick up the pen from the center of the table, then place the pen in the bag”, without a closing period.

一致性:在真实检查点上,LeRobot 策略与 praxis 流水线的 BimanualServeSession 针对相同的三个视图、位姿和指令,通过保存的流水线、在无流水线处理的原始帧上,以及在单个批次中包含两个不同指令的环境时,均产生按位相等的命令(22 个 tick:冷启动块以及深度为 1 的块中的第一条命令)。小型模型测试在 CPU 上针对 praxis Session 重复了该检查。

限制:没有 YAM 机械臂的仿真,因此没有闭环测试套件;forward(微调)仅支持单臂,因为 praxis 仓库目前还没有双臂训练方案。

训练

lerobot-train 使用 praxis 仓库(praxis/train)中的目标函数对策略进行微调:包括视频与动作的联合流匹配损失、文本与本体感觉 dropout、MoE 序列平衡项,以及每次优化器步骤后的路由偏差更新。在相同的窗口和随机抽取下,PraxisPolicy.forward 与 praxis 训练器在已发布检查点上的目标函数在比特级别上完全一致。一个窗口由来自两个相机的 93 个连续帧(分块为 21、24、24、24)组成,包含每一帧的观测法兰位姿和指令法兰位姿以及指令;策略通过 observation_delta_indices / action_delta_indices(93 帧)声明这一点,并使用 drop_n_last_frames(93,以便采样器抽取配方中的窗口起始位置)将窗口保持在剧集内部。相机瓦片遵循检查点的画布规则(video.antialias),与 praxis 数据模块一致。无论启动器如何设置,forward 都在 praxis 数值精度下运行:不使用 autocast(MoE 路由器分数保持 float32),float32 矩阵乘法中不使用 TF32。策略看到的训练窗口数量是一个缓冲区(train_counter):它用于初始化随机抽取,随检查点和适配器一起保存,恢复的运行会继续该序列;保留的评估不会推进它。

数据

该配方在 DROID 上进行训练。lerobot/droid_1.0.1 将笛卡尔位姿存储为六加一列;PR #11 中的 7 维 observation.state.cartesian 和 action.cartesian 列已预先组装好(在合并前请使用 --dataset.revision=refs/pr/11)。数据集自身的 observation.state 和 action 是关节空间的;重命名映射会将笛卡尔列放到相应位置(重命名后的键会替换已有的键),并将两个相机放到策略读取的键下:

RENAME='{"observation.state.cartesian":"observation.state","action.cartesian":"action",
"observation.images.exterior_1_left":"observation.images.exterior","observation.images.wrist_left":"observation.images.wrist"}'

lerobot-train 会将该映射应用于原始批次,并在处理器流水线内部再次应用,因此它不能交换键(例如将 observation.state 映射到另一个名称的操作会被撤销)。

大约一半的 DROID 片段没有指令,且部分片段短于一个窗口;该配方会将它们排除在外。examples/praxis/select_droid_episodes.py 会为 --dataset.episodes 生成列表:

python examples/praxis/select_droid_episodes.py --repo-id lerobot/droid_1.0.1 \
    --revision refs/pr/11 --out droid_episodes.json

Hub 副本存储的是 180x320 的视频;该检查点是在 1280x720 的 DROID 录制数据上训练的。 Hub 副本中的图块是放大后的版本,细节较少,这对于冒烟测试来说已经足够。若要进行忠实的微调,请按照 praxis 仓库的 praxis/train/README.md 中的说明,使用原始视频构建一个本地副本,并将 --dataset.root 指向该副本。

全量微调

权重为 bfloat16;优化器预设 praxis_adamw 是带有随机舍入的 AdamW(普通的 AdamW 会将 3e-6 的更新量舍入掉),调度器采用线性预热和余弦衰减至 3e-7,并且每个 transformer 块在反向传播中都会重新计算其前向过程(activation_checkpointing,默认开启)。权重、梯度和优化器动量每个参数占用 8 字节,对于 27.2B 模型共需 218 GB,并通过 FSDP 在 GPU 之间分配:完整的微调至少需要 4 x 80 GB。praxis 配方为每步 96 个窗口,共 1,200 步;在 8 个 GPU 上,这对应 --batch_size=1 --accelerator.gradient_accumulation.steps=12。lerobot-train 以 micro-batch 为单位计数(一个循环步骤即每个 GPU 上的一个 batch),并按 micro-batch 步进调度器,因此 1,200 次更新对应 --steps=14400,其中预热和衰减均以 micro-step 为单位。

torchrun --nproc-per-node=8 $(which lerobot-train) \
    --dataset.repo_id=lerobot/droid_1.0.1 --dataset.revision=refs/pr/11 \
    --dataset.episodes="$(cat droid_episodes.json)" --rename_map="$RENAME" \
    --policy.path=runwayml/praxis-1-droid-lerobot --policy.device=cuda --policy.push_to_hub=false \
    --parallelism.dp_shard=8 --accelerator.fsdp.mixed_precision_policy.reduce_dtype=float32 \
    --batch_size=1 --accelerator.gradient_accumulation.steps=12 --steps=14400 \
    --policy.scheduler_warmup_steps=600 --policy.scheduler_decay_steps=14400 \
    --num_workers=8 --save_freq=100 --output_dir=outputs/train/praxis_droid

一次保存包含 54 GB 的分片模型(pretrained_model/,可通过 PraxisPolicy.from_pretrained 和 lerobot-eval 加载)以及约 108 GB 的优化器状态。PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True 可避免分配器碎片化,与 praxis 训练器中的做法一致。

在 8 x H100 80 GB 上使用 praxis-1-droid 进行测量,每个 GPU 一个窗口,--dataset.episodes 为 256 个 DROID 剧集(即 Hub 视频),1,000 次优化器步骤:第一步(用于编译内核)之后,每步耗时 3.7 秒,每个 rank 占用 48.5 GB GPU 显存,数据加载器在每个 rank 使用 8 个 worker 的情况下能够跟上进度。损失值在前一百步中从 0.216 降至后一百步的 0.190(视频部分从 0.191 降至 0.171,动作部分从 0.025 降至 0.019)。通过 FSDP,检查点保存分片模型和优化器状态大约需要两分钟。微调后的检查点在 droid_guard 上得分为 50/56,而基础模型为 54:Hub 视频是 180x320 的放大版本,因此在这些视频上进行微调会使策略偏离其训练所用的 1280x720 录制视频。

LoRA

通过 PEFT 使用适配器(pip install -e ".[peft]")。默认目标是 praxis 的 expert 组,以及动作专家的 attention 和 MLP 投影(PraxisPolicy.PEFT_TARGETS);backbone 的 attention 输出投影是第二组,通过 --peft.target_modules 传入。MoE 专家权重不是线性层,因此不添加适配器。在冻结基础权重的情况下,LoRA 训练可以在更少的 GPU 上运行:

torchrun --nproc-per-node=4 $(which lerobot-train) ... --parallelism.dp_shard=4 \
    --peft.r=16 --peft.lora_alpha=32 --output_dir=outputs/train/praxis_lora

保持 --peft.full_training_modules 未设置:该策略的默认 modules_to_save 是训练计数器,适配器必须携带它(配置会拒绝没有它的适配器);它接受的其他模块只有 rig 的三个头部模块,用于从基础检查点进行微调(见下文)。MoE 路由偏差属于基础模型,在 PEFT 下保持固定。在 4 x H100 上,使用 r 16、alpha 32 以及默认目标(--parallelism.dp_shard=4)测量:5.5M 可训练参数,每个 GPU 一个窗口时每步 2.3 秒,每个 rank 30 GB GPU 内存,22 MB 适配器;--parallelism.dp_replicate 配合 dp_shard(HSDP)同样有效。由于冻结的骨干网络负责视频项,总损失在 500 步内保持在 0.22 附近;动作项从 0.026 降至 0.022。该运行生成的适配器在 droid_guard 上得分为 56 分中的 52 分,由 lerobot-eval 加载在基础模型之上。

检查点包含适配器(adapter_config.json、adapter_model.safetensors);lerobot-eval 通过 --policy.path 将其加载到基础模型之上。要为仿真环境获取一个普通的检查点(供服务器和 from_pretrained 使用),请将适配器合并到权重中:

python examples/praxis/merge_praxis_adapter.py --base runwayml/praxis-1-droid-lerobot \
    --adapter outputs/train/praxis_lora/checkpoints/last/pretrained_model --out outputs/praxis-1-droid-lora

与 praxis 训练器的区别

目标、窗口配方、优化器和调度器均采用 praxis 的配置。数据来自 LeRobotDataset,而非 praxis 的文件读取器;随机抽取的随机种子由 LeRobot 种子、调用次数、rank 和 batch 位置共同决定(因此损失值可复现,但与 praxis 运行并非逐步一致);检查点为 LeRobot 检查点。未移植的功能:使用固定抽取的过拟合模式以及 no_sync。

微调基础检查点

praxis-1-base("embodiment": "base",runwayml/praxis-1-base-lerobot)是预训练结束时的视频主干网络和动作专家主干,不包含 rig 的六个头部张量(action_encoder.w1、embed_proprio、head 的权重和偏置),也不包含优化器状态。它不能作为策略加载;有两个脚本可以将其转换为您单臂 rig 微调的起点,即 praxis 仓库中 compute_norm_stats.py 和 init_heads.py 的 LeRobot 版本:

# the normalization statistics of your data: 1 % / 99 % quantiles of the relative actions and the
# mean and std of the proprio over random training windows (no video is read)
python examples/praxis/compute_norm_stats.py --repo-id lerobot/droid_1.0.1 --revision refs/pr/11 \
    --state-key observation.state.cartesian --action-key action.cartesian \
    --episodes droid_episodes.json --n-windows 2000 --out stats.json
# the base plus six new head tensors (`--seed`) and the rig fields: a policy directory that loads
python examples/praxis/init_heads.py --base runwayml/praxis-1-base-lerobot --out outputs/praxis-1-myrig \
    --stats stats.json --tile-height 224 --tile-width 416 --antialias --link
# then the full fine-tune command above with --policy.path=outputs/praxis-1-myrig, or LoRA with the
# six new head tensors trained in full next to the adapters:
torchrun --nproc-per-node=4 $(which lerobot-train) ... --policy.path=outputs/praxis-1-myrig \
    --parallelism.dp_shard=4 --peft.r=16 --peft.lora_alpha=32 \
    --peft.full_training_modules='["train_counter","action_expert.action_encoder.w1","action_expert.embed_proprio","action_expert.head"]'

状态列和动作列保存了 DROID 每帧的位姿和夹爪 [x, y, z, roll, pitch, yaw, gripper](即数据集自身的列,或一个重命名映射)。--stats-from <single-arm checkpoint> 则会复制一个现有的归一化块(即 DROID 机械臂配置)。从基础配置重建的 DROID 机械臂以及已发布检查点的统计量,其配置与 praxis-1-droid 完全一致,且头部张量与 praxis 仓库中针对相同随机种子的 init_heads.py 逐位相同。在上述运行的 DROID 子集上,compute_norm_stats.py 在将夹爪固定为 0 和 1 的情况下,以 0.01 m 的误差复现了检查点的 xyz 分位数统计量。praxis 仓库自身的基础微调(configs/finetune_base.yaml,在 8 块 GPU 上运行一天)从该基础出发,在 droid_guard 上达到了 62.5 %,而使用约 40 倍窗口训练的 praxis-1-droid 则达到了 91 %:该基础配置是新机械臂的教程起点,而非已发布检查点的替代品。完整微调会在 optimizer_lr 下训练所有权重;在 LoRA 下,仅适配器和六个头部张量参与训练(PraxisPolicy.PEFT_HEAD_MODULES 是除计数器外 full_training_modules 接受的唯一模块),主干的其余部分保持冻结;praxis 配方中的仅头部预热步骤和头部学习率缩放并未移植过来。

已知限制

  • 两套机械臂:单臂 DROID 检查点(embodiment: single_arm)和双臂 YAM 机械臂(bimanual)。微调与仿真均为单臂。
  • Praxis 为配备阻抗控制器的机械臂输出笛卡尔位姿目标;LeRobot 的机器人类中不包含 Franka,因此实际机器人的使用需要通过一个执行笛卡尔目标的客户端来实现。
  • select_action 批次会按顺序运行其环境(每个环境调用一次模型)。
  • 不支持远程服务。默认的 chunk runner(lerobot/inference/policy_runner.py)会拒绝依赖于过去观测的策略,而 Praxis 依赖于已执行 chunk 的帧(在 chunk_inference_spec 中 current_observation_only=False);tile 步骤还会将相机帧调整到检查点画布的大小,因此它们与 validate_chunk_input_features 检查的形状不匹配。请在进程内运行该策略(lerobot-eval,一个机器人客户端)。
Update on GitHub