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 は、Runway のロボット向けオープンウェイト World Action Model である Praxis-1 の LeRobot アダプターです。Praxis-1 はビデオアクションモデルであり、チャンク因果的なビデオ拡散バックボーン(27B パラメータの mixture-of-experts トランスフォーマーで、トークンあたり約 3B がアクティブ)が、2 つのリグカメラの次の 1.6 秒間のビデオをデノイズする一方で、バックボーンの隠れ状態を読み取るアクションエキスパートが、対応する 24 個のデカルトコマンドを 15 Hz でデノイズします。Wan2.2 ビデオオートエンコーダーと Qwen3.5-2B テキストエンコーダーは凍結されており、それぞれの公開ウェイトから提供されています。公開されたチェックポイント praxis-1-droid は、単一アームの DROID リグを駆動します。これは、外部カメラと手首カメラを備えた Franka アームです。以前の未公開の DROID チェックポイントは、他のウェイトと別のカメラキャンバスを持ちますが、同じアーキテクチャです。チェックポイントの config.json がそれらを区別します(「The checkpoints」を参照)。2 つの YAM アーム用のバイマニピュレーションチェックポイントである praxis-1-yam は、同じレイアウトで提供されています(「The bimanual YAM checkpoint」を参照)。リグのヘッドを含まないビデオバックボーンとアクションエキスパートのトランクである praxis-1-base は、新しい単一アームリグのファインチューニングの出発点です(「Fine-tune the base checkpoint」を参照)。

リソース:

リソースリンク用途
アナウンス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凍結

モデルコードは、praxis リポジトリのコミット 8dad1a4 からベンダー化し、praxis-1-droid の canvas ルールを追加してインポートパスを書き換えたうえで、lerobot/policies/praxis/praxis_core 以下にツリー内に配置されています(ファイルヘッダーにアップストリームのコミットが明記されています)。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: CUDA 13.2 付き torch 2.15.0、transformers 5.22.2、triton 3.9)で実行されます。以下の結果は、そのスタック(transformers 5.22.1 と 5.22.2 は同じコマンドを生成します)および LeRobot 0.6 のスタック(CUDA 12.8 付き torch 2.11.0、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 実装を実行します。これら 3 つはモデルではなく、出力の最後のビットを変更します。

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が1つ必要です(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 トークンのブロック単位でのみ受け付けます(ValueError: Block sparse tensorsforward have block size 128, which must be a multiple of 256 が発生します)。 flash バックエンドでは、これらの GPU 上で、チェックポイントの backbone.attention.flex_block_size の 128 ではなく、256 トークンのブロック単位でバックボーンのマスクを構築します。ブロックサイズはトークンペアをグループ化するだけです。マスク関数は依然として各ペアを決定します。praxis リポジトリでは、チェックポイントの config.json に "flex_block_size": 256 を設定して同じように実行します。triton バックエンドは、チェックポイントの 128 を維持します。

RTX Blackwell GPU(compute capability 12.x。RTX PRO 6000 Blackwell は、重みの必要とする 96 GB を備えています)では、FlashAttention-4 はどのブロックサイズでも block-sparse forward をサポートしていません(チェックポイントの 128 トークンブロックでは Block sparsity requires sparse_block_size[1]=128 to match tile_n=64 が発生します)。そのため、flash バックエンドはこれらの GPU で実行できません。attention_backend="auto" は、flash-attn-4 がインストールされている場合でも、これらの GPU では triton を選択し、"flash" を指定するとポリシー構築時に ValueError が発生します。praxis リポジトリは、チェックポイントの flash を使用すると最初の呼び出しで失敗しますが、チェックポイントの config.json に "backend": "triton" を設定すると同じように動作します。2 つのバックエンドは最後のビットで異なるため、これらの GPU での実行結果は、他の場所での triton 実行結果と比較してください。RTX PRO 6000 において、それ以外の変更はありません。以下の 2 つのスタックの 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 を読み込む際、保存されている model.language_model.* という名前を照合する前にリネームするため、言語モデルのテンソル 320 個すべてが欠損していると報告され、無言のうちにランダムに初期化されてしまいます。ポリシーのテキストエンコーダーはこれを検出し、ログに警告を出力しながら、テンソルを保存時の名前で読み込みます。テキストエンコーダーがランダムなままの状態で残ると、ポリシーは動き続けるものの、成功することはほとんどありません(最初のランでは、ガードエピソード 56 回中 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フランジの姿勢(メートル、ラジアン単位の外積x-y-zオイラー角、DROID規約)および[0, 1]のグリップ
taskstrテキストエンコーダーにそのまま渡される指示。
action(7,)同じ7つの値:次の15 Hzティックの絶対ターゲットポーズとグリップ
observation.scene_seed(任意)(1,)エピソードのシーンシードです。チャンクkのサンプラーシードはseed + seed_stride * scene_seed + kです。

ポリシーは独自の履歴を保持します。バッチ内の各環境について、直近の25回の観測(実行されたチャンクと現在のフレーム)をリングバッファとして保持し、これを次の呼び出しのコンテキストとして使用します。そのため、n_obs_stepsは1のままとなり、select_actionは15 Hzの各ティックで、そのティックの観測を引数として呼び出す必要があります。コンテキストウィンドウの最初の呼び出しでは21個のコマンドが返され、それ以降の呼び出しでは24個が返されます。n_action_stepsが24未満の場合、praxisと同様に、すべての呼び出しがコンテキストなしのコールドスタートとなります。drop_queued_actions(ロールアウト中の新しい指示など)も、モデルが実行された直前のチャンク全体に条件付けられているため、時間的なコンテキストをリセットします。チャンク全体を実行する呼び出し元(predict_action_chunk)は、前回の呼び出し以降に実行されたフレームをスタックとして引き継ぎます。これは、古い順に並べた(B, T, 3, H, W)のカメラと(B, T, 7)の状態です。次のチャンクは、ティックごとに呼び出すロールアウトのものと同一になります。1フレームでのチャンク呼び出しは、コンテキストを再度開始します。

プロセッサ

preprocessorは、praxisのデータおよびサーブコードが保存した形式で、検査可能なステップとして、入力を準備します。

ステップDoes
PraxisTileViewsProcessorStep各カメラビューは、チェックポイントのルール(video.antialias)に従ってキャンバスのタイルになります。praxis-1-droid: ネイティブの 720 x 1280 アスペクト比にセンタークロップし、234 x 416 にアンチエイリアス付きバイリニアリサイズを 1 回行います(そのサイズで送信されたビューの場合は行いません)、その後 224 x 416 にセンタートリムします。以前の DROID チェックポイント: タイルのアスペクト比にセンタークロップし、256 x 416 に通常のバイリニアリサイズを 1 回行います。uint8 レベル; タイルサイズに既に一致しているビューはそのまま通過します。
PraxisPoseProcessorStepDROIDのポーズ[x, y, z, roll, pitch, yaw, gripper]は、praxisのポーズ8[x, y, z, qw, qx, qy, qz, gripper](状態、およびトレーニングバッチのアクション)になります。
PraxisTokenizerProcessorStep指示は、Qwen3.5-2B テキストエンコーダーのトークン ID とアテンションマスク(observation.language.tokens、.attention_mask)に変換されます。トークナイゼーションは、praxis テキストエンコーダーが行うものと同じ方法で行われます。凍結されたエンコーダーとその埋め込みキャッシュは、ポリシー内に保持されます。
PraxisCommandProcessorStep(postprocessor) ポリシーのコマンド、praxis ワイアースの姿勢 [x, y, z, rotation 3 x 3] およびグリップが、CPU 上で DROID の 7 値に変換されます。

これらを通じて、ポリシーはタイルと pose8 行を認識し、ワイヤーポーズを出力します。これはまさに praxis の Session が消費し、生成するものです。q99 アクション正規化、平均/標準偏差によるプロプライオ正規化、および相対アクションのデコードはセッション内に保持されます。これらは各実行チャンクの開始時のポーズにアンカーされており、セッションはロールアウト全体にわたってこれを追跡します。また、その統計量はチェックポイントの PraxisConfig.model_config["normalization"] です。標準の正規化ステップは恒等変換です。ポリシーはまた、生の 7 値のポーズ、生のカメラフレーム、および task 文字列も受け入れるため、パイプラインなしで呼び出すことができます。タイルサイズのフレームはタイルとして扱われ、その他のサイズはチェックポイントのキャンバスルールに従います(praxis-1-droid のフィットルールは冪等ではないため、タイルが 2 回通ることはありません)。トークンなしのタスクはポリシーによってトークン化されます。タイルはオートエンコーダーが読み取るまで、パイプラインが配置したデバイス上に保持され、コマンドは CPU に返されます。デバイス境界を 2 回またぐことはありません。

チェックポイント

チェックポイントタイルリサイズ推奨サンプラー(praxis-1-droidのPraxisConfigのデフォルト値)
praxis-1-droid(リリース済み)224 x 416アンチエイリアス適用ルール(video.antialias: true)2回のEulerステップ、flow shift 12.5、テキストガイダンスは動画に4.0、アクションにはなし(1.0)、1つのコンテキストチャンク
YAM、praxis-1-yam(両手操作、リリース済み)256 x 416、3つのタイルYAM fit ルール(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 416通常の双線形リサイズ2ステップ、シフト8.0、テキストガイダンスなし、3つのコンテキストチャンク(praxis)。以下の結果は、シフト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から取得されます。両手操作のチェックポイントは、3つのカメラキー(--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 のシャードを書き出します)。また、2 つのプロセッサファイルも含まれます。--vae-id と --text-encoder-id は、フリーズされた部分のソースを指定します(デフォルトでは Hub の ID です)。

結果

以下のクローズドループの数値は、praxis リポジトリの MuJoCo リグ(praxis-sim パッケージ: Franka + Robotiq 2F-85、512 x 288 カメラ、15 Hz)およびその固定評価スイートから得られたものです。エピソードの LeRobot シードは、そのシーンのシードです。lerobot-eval によってスイートを実行する環境は、LeRobot の一部ではありません。

サンプリング中、ポリシーはpraxisが実行に使用する数値設定(cuDNNのベンチマークモードなし、MoEルーター用のfloat32行列積、cuDNNでのTF32許可)を固定し、その後ホスト設定を復元します。そのため、--tf32、torch.backends.cudnn.benchmark、およびバッチサイズは出力を変更しません。

droid_guardスイート(7つのDROIDタスク x 8シーン、56エピソード、サンプラーシード0)を、以前の未公開DROIDチェックポイント(256 x 416タイル)で実行しました。サンプラー設定は、当時のpraxisリポジトリのものです(Eulerステップ2、flow shift 1.0、action guidance 3.0、video guidanceなし、context chunk 1)。また、flashアテンションバックエンドを使用し、H100 80 GB 1枚で実行しました(約10分)。同じエピソードを5つの方法で実行しました。praxisリポジトリ自身のランナーを自身の環境(torch 2.13、transformers 5.2.0、causal-conv1d)で実行、同じランナーとlerobot-evalをLeRobot 0.6のスタック(torch 2.11、transformers 5.5.4)で実行、そして両方をLeRobot 0.7が固定するスタック(CUDA 13.2付きtorch 2.15、triton 3.9、flash-linear-attention 0.6.0、当時のtransformers 5.16.1)で実行しました。torch 2.15の実行をtransformers 5.22.1で繰り返すと、praxisランナーと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パイプラインとビット単位で等しくなります(下記を参照)。2つのスタック間では、同一の入力に対するコマンドの差は最大0.04(flash)および0.02(triton)であり、これが個々のエピソードの挙動を左右します。

公開されたチェックポイントであるpraxis-1-droid

praxis-1-droid(224 x 416 のアンチエイリアス処理済みタイル)の推奨設定(Euler ステップ 2、flow shift 12.5、video guidance 4.0、action guidance なし、context chunk 1)で、同じ 56 エピソードを 1 台の H100 80 GB 上で flash バックエンドを使用して実行した結果です。具体的には、praxis リポジトリのランナー(praxis 環境:torch 2.13、transformers 5.2.0 における現在のサンプラー付きの main)、LeRobot 0.6 のスタック(torch 2.11、transformers 5.5.4)での lerobot-eval、そして LeRobot 0.7 が固定するスタック(CUDA 13.2 付きの torch 2.15、transformers 5.22.1)でのランナーおよび lerobot-eval の 3 つです。

タスク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ランナーを引いた差は、両方のスタックで+0.0ポイントでした(95%ブートストラップ区間 [-7.1 %, +7.1 %]、2勝2敗、McNemar p 1.00)。結果が反転するエピソードは、2つのスタックで同じものではありません。torch 2.15のカーネルがコマンドを最大0.02移動させるため、境界付近のいくつかのエピソードの結果がそれに伴って変化します。praxis-1-droidとその設定ではpress_buttonを解決できますが、以前のDROIDチェックポイントとその設定では、主に失敗していました。

praxisリポジトリとの互換性

  • 1つのプロセス内で同一の観測に対して、LeRobotのポリシーとpraxisのSessionは、ビット単位で等しいテキスト埋め込み、VAE潜在変数、プロプライオ、サンプリングされたアクションおよびコマンドを出力します(読み込まれた重みが同一である場合)。tests/policies/praxis/test_praxis.pyでは、CPU上で小さなランダムモデルを用いてこのチェックを繰り返し実行します。
  • プロセス間で、元のpraxisパッケージとLeRobotのポリシーが実行するpraxisパイプラインは、実際のチェックポイントに対して、flashバックエンドでもtritonバックエンドでも(2つのバックエンド間のコマンドの差は最大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)の両方で確認されています。参照値自体は2つのスタック間で変動します(flashでは最大0.019、tritonでは0.015)。ポリシーは、それぞれのスタックで参照値をビット単位で追従します。
  • praxis-1-droid: praxis リポジトリの main(現在のサンプラー付き)の praxis パイプラインに対するプロセス間チェックは、ポリシーを呼び出す 8 つの方法において、両方のバックエンドおよび両方のスタックでビット単位で等しいことが確認されました: 保存されたプロセッサのパイプラインを介して(lerobot-eval が行うように)、生フレームと 7 値のポーズに対して パイプラインなしで、1 つのバッチに 2 つの環境(異なる指示とシーンのシード)を含めて、predict_action_chunk を介して(コールドスタート、および実行されたフレームをスタックとして与えた次のチャンク)、n_action_steps=12(12 コマンドごとにチャンク、実行されたフレームに条件付け)、アクションガイダンスをオン(cfg_scale=3)、および 3 つのコンテキストチャンク(深度 0 から 3)で 70 コマンドにわたって。このチェックは、lerobot-eval のデフォルトの TF32 および cuDNN ベンチマークフラグの下でも成立します: ポリシーはサンプリング中に praxis の数値を固定し、ホストにそのフラグを返します。タイルステップは、フレームのデバイスに関係なく、キャンバスを CPU で計算します: CUDA のアンチエイリアスリサイズは、約 3 分の 1 のピクセルで 1 レベル異なる丸めを行い、ステップがそうする前にコマンドを最大 0.02 移動させていました。
  • Training: 公開された重みにおいて、praxis リポジトリの PraxisPolicy.forward と praxis.train.objective は、同じウィンドウと乱数抽出に対して、ビット単位で同じ損失(および動画とアクションの項)を出力し、アクションエキスパートの全 200 個のテンソルに対して同じ勾配を出力します。報告されている MoE バランス指標は、どちらの側でも実行ごとに 1〜2 個の 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 トークンのマスクブロックを使用)でビット単位で等しい結果が得られました。具体的には、praxis-1-droid についてはポリシーを呼び出す 7 つの方法(lerobot-eval のフラグで保存されたパイプライン、生の 720 x 1280 フレーム、1 バッチで 2 つの環境、フレームスタック付きの predict_action_chunk、n_action_steps=12、cfg_scale=3、および 3 つのコンテキストチャンクを持つ 94 コマンド)、praxis-1-yam についてはパイプライン経由で生フレームおよび 2 つの環境を使用した場合、そして 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 と同様にビット単位で等しい: 7 つ (463 コマンド、コンテキスト深度 0 から 3) で praxis-1-droid、3 つで praxis-1-yam、そして praxis-1-base から再構築された DROID リグ (その 6 つのヘッドテンソルと設定は、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 のスタックで droid_guard における praxis-1-droid によるクローズドループ: 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、デフォルトのターゲット、ステップごとに 1 ウィンドウ) は 1 つの 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 に対するプロセス間チェックは、GB300 と同様に、両方のバックエンド(flash は両側で 256 トークンのマスクブロックを使用)でビット単位で等しくなります。具体的には、praxis-1-droid は 7 つ(バックエンドおよびスタックごとに 463 コマンド、コンテキスト深度 0 から 3)、praxis-1-yam は 3 つ、そして praxis-1-base から再構築された DROID リグ(その 6 つのヘッドテンソルと設定は、praxis リポジトリの init_heads と等しい)です。2 つのバックエンドは、コマンドにおいて最大 0.19 だけ異なります。forward は、praxis の目的関数の損失、その動画およびアクションの項、そして学習済みバックボーンにおけるすべての 200 個のアクションエキスパートテンソルの勾配を、両方のバックエンドおよび両方のスタックでビット単位で一致させます。lerobot-train による LoRA ファインチューニングは 1 つの GPU で実行可能であり(ステップあたり 1.2 秒、62 GB)、アダプターとそのマージ済みチェックポイントは lerobot-eval で読み込むことができます。
  • B200 でのクローズドループ(droid_guard、praxis-1-droid、flash): praxis ランナーは 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 ランナーはセッションのポーズとコマンドをそのまま渡します。この往復でポーズは約 1e-7 移動し、エピソード全体を通じてその差は別の軌跡にまで拡大します。praxis ランナーが同じ float32 変換を経由して状態とコマンドを送信すると、両方のスタックで全 56 エピソードにおいて lerobot-eval と同じ結果になります。
  • praxis リポジトリ自身の環境(torch 2.13.0+cu130、transformers 5.2.0、causal-conv1d)に対して、出力はビット単位で等しくありません。カーネルが異なり、Qwen3.5 の埋め込みは bfloat16 で最大 0.125 まで異なり、アクションの最後のビットもそれに従います。したがって、リポジトリ間のチェックは、上記のペアリングされたクローズドループ評価です。
  • シミュレーションは4xマルチサンプリングでレンダリングされ、同じシーンのレンダリング間でフレームごとに少数のピクセルが異なるため、同じポリシーのクローズドループエピソードは、重みが同一であっても最初のチャンク以降で分岐します。ペアのエピソードはフレーム単位ではなく統計的に比較されます。

両手YAMチェックポイント

praxis-1-yam("embodiment": "bimanual"、runwayml/praxis-1-yam-lerobot)は、トップカメラと2つのリストカメラ(16:10 の ZED ビュー、1920 x 1200)を備えた2つの YAM アームを駆動します。同じモデルが実行されますが、エキスパートヘッドの幅(アームごとに20のアクション値: [xyz, rot6d, gripper])、3タイルのキャンバス、正規化がチェックポイント間で異なります。LeRobot アダプターは、praxis リポジトリの praxis/bimanual.py(praxis_core/bimanual.py としてベンダー化)に従います。

キー形状意味
observation.images.top_left, observation.images.wrist_left, observation.images.wrist_rightRGBキャンバス順の3つのビュー(camera_keys)。16:10の任意のサイズ。プロセッサはYAMフィットルール(ビュー全体が416 pxのボックスに収まるようにし、その後32 pxトリムする)で256 x 416のタイルを作成します。
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.

Parity: 実際のチェックポイントにおいて、LeRobot ポリシーと praxis パイプラインのBimanualServeSessionは、同じ3つのビュー、ポーズ、指示に対して、保存されたパイプラインを介して、それらを介さずに生フレームに対して、そして1つのバッチ内で異なる指示を持つ2つの環境に対して、ビット単位で等しいコマンドを出力します(22ティック: コールドスタートチャンクと深度1チャンクのコマンドの先頭)。tiny-model テストでは、praxis のSessionに対して CPU 上でこのチェックを繰り返します。

制限事項: YAM リグのシミュレーションが存在しないため、クローズドループのスートはありません。forward(ファインチューニング)は、praxis リポジトリにバイマヌアル(両手)のトレーニングレシピがまだないため、シングルアームのみ対応しています。

トレーニング

lerobot-train は、praxis リポジトリ(praxis/train)の目的関数を目標としてポリシーをファインチューニングします。具体的には、動画とアクションのジョイント・フローマッチング損失、テキストおよびプロプライオのドロップアウト、MoE のシーケンスバランス項、そして各オプティマイザーステップ後のルーティングバイアス更新です。同じウィンドウと乱数抽選が与えられた場合、PraxisPolicy.forward は、公開されたチェックポイント上で praxis トレーナーの目的関数とビット単位で等しくなります。ウィンドウとは、両方のカメラの 93 連続フレーム(21、24、24、24 のチャンク)で、各フレームの観測されたフランジ姿勢と指示されたフランジ姿勢、および指示文を含みます。ポリシーは observation_delta_indices / action_delta_indices(93 フレーム)でこれを宣言し、drop_n_last_frames(93、これによりサンプラーがレシピのウィンドウ開始位置を抽選します)を用いてエピソード内にウィンドウを保持します。カメラのタイルは、praxis データモジュールと同様に、チェックポイントのキャンバスルール(video.antialias)に従います。forward は、ランチャーが何を設定しても、praxis の数値環境下で実行されます:オートキャストなし(MoE ルーターのスコアは float32 のまま)、float32 の行列積では TF32 なし。確認されたトレーニングウィンドウの数は、ポリシーのバッファ(train_counter)です。これは抽選のシードとなり、チェックポイントとアダプターとともに移動し、再開された実行ではシーケンスが継続されます。ホールドアウト評価では、このカウンターは進みません。

データ

このレシピは DROID で学習を行います。lerobot/droid_1.0.1 は、Cartesian ポーズを 6 列 + 1 列として保持しています。PR #11 の 7 次元の observation.state.cartesian および action.cartesian 列はこれらを事前に取り揃えています(マージされるまで --dataset.revision=refs/pr/11 を使用してください)。データセット自体の observation.state および action はジョイント空間です。リネームマップは、Cartesian 列をその位置に配置し(リネームされたキーは既存のキーを置き換えます)、ポリシーが読み取るキーの下に 2 つのカメラを配置します:

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 レシピは 1 ステップあたり 96 ウィンドウ、1,200 ステップです。8 GPU での実行では --batch_size=1 --accelerator.gradient_accumulation.steps=12 となります。lerobot-train はマイクロバッチをカウントし(1 ループステップはすべての GPU での 1 バッチ)、スケジューラーをマイクロバッチごとにステップします。そのため、1,200 回の更新は --steps=14400 となり、ウォームアップと減衰はマイクロステップ単位で設定されます。

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

1つの保存には、54 GBのシャーディング済みモデル(pretrained_model/、PraxisPolicy.from_pretrainedおよびlerobot-evalで読み込めます)と、約108 GBのオプティマイザ状態が含まれます。PYTORCH_CUDA_ALLOC_CONF=expandable_segments:Trueにより、praxis trainerと同様にアロケータの断片化を回避できます。

8 x H100 80 GB 上で、praxis-1-droid を使用し、GPU ごとに 1 ウィンドウ、--dataset.episodes を 256 の DROID エピソード(Hub の動画)に設定し、1,000 ステップのオプティマイザ実行を行った結果、最初のステップ(カーネルをコンパイルする)以降は 1 ステップあたり 3.7 秒、ランクあたり GPU メモリ 48.5 GB、ランクあたり 8 ワーカーでデータローダーが追従できる状態でした。損失は最初の 100 ステップで 0.216 から、最後の 100 ステップで 0.190 まで低下しました(動画は 0.191 から 0.171、アクションは 0.025 から 0.019)。チェックポイントは、FSDP を通じてシャードされたモデルとオプティマイザの状態を約 2 分で保存します。ファインチューニングされたチェックポイントは droid_guard で 56 点中 50 点、ベースモデルは 54 点でした。Hub の動画は 180x320 のアップスケールであるため、これらでのファインチューニングは、1280x720 の録画で学習されたポリシーから逸れてしまいます。

LoRA

PEFT によるアダプター(pip install -e ".[peft]")。デフォルトのターゲットは、praxis の expert グループと、アクションエキスパートの attention および MLP の射影(PraxisPolicy.PEFT_TARGETS)です。バックボーンの attention 出力射影は第 2 グループであり、--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 は学習カウンターであり、アダプターが保持する必要があります(この設定では、それを含まないアダプターは拒否されます)。それ以外に受け入れられるモジュールは、ベースチェックポイントからのファインチューニング(下記参照)のために、リグの3つのヘッドモジュールのみです。MoEルーティングバイアスはベースに属し、PEFT下では固定されます。r 16、alpha 32、デフォルトのターゲットで 4 x H100(--parallelism.dp_shard=4)で測定したところ、学習可能なパラメータは 5.5M、GPUあたり1ウィンドウでステップあたり 2.3 秒、ランクあたり GPU メモリ 30 GB、アダプターは 22 MB でした。--parallelism.dp_replicate と dp_shard(HSDP)の組み合わせも動作します。凍結されたバックボーンが動画項を支配するため、500ステップにわたって総損失は 0.22 付近に留まります。アクション項は 0.026 から 0.022 へ変化しました。その実行のアダプターは、lerobot-eval によってベース上に読み込まれた droid_guard で 56 点中 52 点を記録しました。

チェックポイントにはアダプター(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 trainerとの違い

目的、ウィンドウレシピ、オプティマイザー、スケジュールは praxis のものです。データは praxis のファイルリーダーではなく LeRobotDataset から取得され、乱数抽出は LeRobot のシード、呼び出し回数、ランク、バッチ位置からシード付けされます(そのため損失は再現可能ですが、praxis の実行とステップ単位で一致するわけではありません)。また、チェックポイントは LeRobot のチェックポイントです。移植されていないもの: 固定の乱数抽出と no_sync を伴う過学習モード。

ベースチェックポイントをファインチューニングする

praxis-1-base("embodiment": "base"、runwayml/praxis-1-base-lerobot)は、事前学習の最終段階における動画バックボーンおよびアクションエキスパートのトランクです。リグの6つのヘッドテンソル(action_encoder.w1、embed_proprio、headの重みとバイアス)やオプティマイザの状態は含まれていません。ポリシーとして読み込むことはできず、単一アームのリグ向けのファインチューニングの開始点に変換するには、praxisリポジトリのcompute_norm_stats.pyとinit_heads.pyのLeRobot版である2つのスクリプトを使用します。

# 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"]'

state 列と action 列には、各フレームごとの 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 に固定した状態で、xyz の分位数においてチェックポイントの統計情報を 0.01 m 以内で再現します。praxis リポジトリ独自のベースファインチューニング(configs/finetune_base.yaml、8 GPU で 1 日)は、このベースから droid_guard で 62.5 % を達成しますが、約 40 倍のウィンドウ数で学習した praxis-1-droid は 91 % です。ベースは新しいリグのチュートリアル用の出発点であり、公開されたチェックポイントの代替ではありません。完全なファインチューニングでは、すべての重みが optimizer_lr で学習されます。LoRA の場合、アダプタと 6 つのヘッドテンソルが学習対象となり(PraxisPolicy.PEFT_HEAD_MODULES は、full_training_modules が受け入れるモジュールのうち、カウンターを除いて唯一のものです)、トランクの残りは凍結されます。praxis レシピのヘッド専用のウォームアップステップとヘッドの学習率スケールは移植されていません。

既知の制限

  • 2つのリグ:単腕のDROIDチェックポイント(embodiment: single_arm)と、両腕のYAM(bimanual)です。ファインチューニングとシミュレーションは単腕です。
  • Praxisはインピーダンスコントローラを備えたアームに対してデカルト座標の姿勢目標を出力します。LeRobotのロボットクラスにはFrankaが含まれていないため、実機での利用はデカルト座標の目標を実行するクライアントを介して行われます。
  • select_action バッチは、環境を順番に実行します(環境ごとに1回のモデル呼び出し)。
  • リモートサービングはサポートされていません。デフォルトのチャンクランナー(lerobot/inference/policy_runner.py)は、過去の観測に条件付けするポリシーを拒否しますが、Praxis は実行されたチャンクのフレームに条件付けします(chunk_inference_spec 内の current_observation_only=False)。また、タイルステップではカメラフレームがチェックポイントのキャンバスにリサイズされるため、validate_chunk_input_features がチェックする形状と一致しなくなります。ポリシーをインプロセスで実行してください(lerobot-eval、ロボットクライアント)。
Update on GitHub