# LingBot-VA

LingBot-VAは、**Wan2.2**のビデオ拡散スタックを基盤とする**自己回帰型ビデオ・アクション・ワールドモデルポリシー**です。1つの自己回帰シーケンス内で、将来の**ビデオ潜在変数**と**ロボットアクション**の予測を交互に行います（「VA」= Video-Action）。LeRobotの統合により、LingBot-VAは標準的な学習、評価、プロセッサのインターフェースに接続されます。

## モデル概要

LingBot-VA is a **dual-stream "mixture-of-transformers"**: a video/latent stream (`patch_embedding_mlp → blocks → proj_out`) and an action stream (`action_embedder → blocks → action_proj_out`) share the same 30 transformer blocks and text conditioning.

| コンポーネント                | クラス                   | 役割                                                        |
| ------------------------ | ----------------------- | ----------------------------------------------------------- |
| DiTバックボーン（学習可能） | `WanTransformer3DModel` | 約50億パラメータのデュアルストリームTransformer。                          |
| VAE（凍結）             | `AutoencoderKLWan`      | Wan2.2 VAE、`z_dim=48`。ソースリポジトリから遅延取得されます。   |
| テキストエンコーダー（凍結済み）    | `UMT5EncoderModel`      | UMT5-XXL、`d_model=4096`。ソースリポジトリから遅延取得されます。 |

推論時、ポリシーはチャンクごとに自己回帰ループを実行します。2つの独立したフローマッチングスケジューラを用いて、ビデオ潜在ストリーム（CFG、約20ステップ）とアクションストリーム（約50ステップ）をデノイズし、チャンク間でKVキャッシュを維持します。チャンクが実行される間に、実際に観測されたキーフレームがKVキャッシュにフィードバックされます（クローズドループのワールドモデリング）。

### LeRobot統合の範囲

- LeRobot による標準的な `policy.type=lingbot_va` 設定。
- Hub 上の LeRobot 形式のチェックポイント（公開されたアップストリームのものから変換済み）
- 標準的な`select_action`インターフェースの背後にある、自己回帰型デュアルストリーム推論（単一環境の評価、`--eval.batch_size=1`）。
- eval / training 中に、ポリシーの**予測（想像）された動画**を保存する機能をオプトインで有効にします。
- LIBEROとRoboTwinでの`lerobot-eval`による評価。
- `policy.forward` を用いたデュアルストリームフローマッチング損失による学習 / 微調整（詳細は下記を参照）。

## インストール

1. [インストールガイド](./installation)に従ってLeRobotをインストールしてください。
2. LingBot-VAの追加コンポーネントをインストールします:

```bash
pip install -e ".[lingbot_va]"
```

## チェックポイント

リリースされたアップストリームのチェックポイントはLeRobot形式に変換され、Hubにプッシュされました:

| バリアント                | LeRobotチェックポイント               |
| ---------------------- | -------------------------------- |
| LIBERO-Long 後学習 | `lerobot/lingbot_va_libero_long` |
| RoboTwinのポストトレーニング    | `lerobot/lingbot_va_robotwin`    |
| 事前学習済みベース        | `lerobot/lingbot_va_base`        |

学習可能な ~5B の Transformer のみが LeRobot の `model.safetensors` に保存されます。凍結された VAE + UMT5 + トークナイザー（約 20 GB）は、読み込み時に `config.wan_pretrained_path` から取得されます（デフォルトではソースの `robbyant/*` リポジトリを参照します）。UMT5-XXL テキストエンコーダーはデフォルトで CPU 上で実行されます（`config.text_encoder_device`）ため、5B の Transformer と VAE が単一の 24–32 GB の GPU に収まります。

## Evaluation (LIBERO)

```bash
lerobot-eval \
    --policy.path=lerobot/lingbot_va_libero_long \
    --policy.device=cuda \
    --env.type=libero --env.task=libero_10 \
    --env.observation_height=128 --env.observation_width=128 \
    --eval.n_episodes=50 --eval.batch_size=1 \
    --output_dir=outputs/eval/lingbot_va_libero
```

LingBot-VAのストリーミング推論（KVキャッシュ＋観測キーフレームのフィードバック）は単一環境の評価用に実装されています。`--eval.batch_size=1`を使用してください。

## Evaluation (RoboTwin)

RoboTwin 2.0にはSAPIEN + CuRoboシミュレータースタックが必要です。`docker/Dockerfile.sim`の`robotwin_native`ステージは参照用のインストール手順です（LingBot-VAも`warp-lang==1.3.1`と固定された`pytorch3d`、および`TORCH_CUDA_ARCH_LIST`にGPUの計算能力を指定してビルドしたCuRoboが必要です）。シミュレータサーバーのRoboTwin設定はジョイント制御を使用しているため、以下のエンドエフェクタモードはローカルで実行されます。RoboTwinは**エンドエフェクタ姿勢制御**を使用するため、`--env.action_mode=ee`で実行してください。ポリシーはアームごとの`xyz+quaternion+gripper`のデルタ（`robotwin_tshape`潜在レイアウト）を予測し、エピソードの初期eef姿勢に合成してCuRoboのIKで実行します。

```bash
lerobot-eval \
    --policy.path=lerobot/lingbot_va_robotwin \
    --policy.device=cuda \
    --env.type=robotwin --env.task=beat_block_hammer --env.action_mode=ee \
    --eval.n_episodes=10 --eval.batch_size=1 \
    --output_dir=outputs/eval/lingbot_va_robotwin
```

### 予測（想像）された動画の保存

`--policy.save_predicted_video=true` を設定すると、予測された動画の潜在変数も VAE でデコードされ、環境レンダリングされた `eval_episode_*.mp4` ビデオの隣に `pred_episode_*.mp4` が書き出されます。このフラグは、`lerobot-train` 実行中の定期評価にも同様に機能します。

### 想像上の動画をライブで視聴する

`lerobot-rollout --display_data=true` を実行する際に `--policy.return_predicted_video=true` を設定すると、各チャンクの予測動画をデコードし、各アクションの実行に伴って、それに属する想像上のカメラフレーム（各実カメラの隣に `prediction.images.<camera>` として表示）を表示します。これにはチャンクごとに1回のVAEデコードのコストがかかります。

## トレーニング / ファインチューニング

`LingBotVAPolicy.forward(batch)` は、論文に記述されているデュアルストリームの**flow-matching**損失（`latent_loss + action_loss`、タイムステップ重み付き、アクションマスク付き）を実装しています。具体的には、カメラクリップをVAEでエンコードしてビデオ潜在変数に変換し、タスクをUMT5でエンコードし、両方のストリームにノイズを加えた後、Transformerのブロック因果学習パスを実行し、`(loss, metrics)`を返します。オプティマイザーのプリセットはAdamWで、線形ウォームアップ後に一定になるスケジュールを使用しています（アップストリームと一致）。

要件:

- ブロック因果マスクは PyTorch の **flex-attention** を使用するため、学習には `--policy.attn_mode=flex` でポリシーを構築してください（デフォルトの `torch` SDPA は推論専用です）。
- 完全な5B DiTは、AdamWを使用する場合、単一の24〜32 GBのGPUに収まりません。**LoRA**（`--peft.method_type=LORA`）および/またはオプティマイザーのオフロードを使用してファインチューニングしてください。`get_optim_params`は、学習可能なパラメータ（アダプターなど）のみを返します。VAEとUMT5テキストエンコーダーは凍結されたままです。PEFTサポートを有効にするには、`lerobot[peft]`エクストラをインストールしてください（[PEFTトレーニングガイド](./peft_training)を参照）。

```bash
lerobot-train \
  --policy.path=lerobot/lingbot_va_libero_long --policy.attn_mode=flex \
  --peft.method_type=LORA --peft.r=32 --peft.lora_alpha=32 \
  --peft.target_modules='transformer\.blocks\.\d+\.attn[12]\.(to_q|to_v)' \
  --dataset.repo_id=<your LeRobot-format dataset> \
  --batch_size=1 --steps=... --output_dir=outputs/train/lingbot_va
```

SmolVLA / π₀ とは異なり、LingBot-VA は組み込みのデフォルト LoRA ターゲットを同梱していないため、`--peft.target_modules` を明示的に指定する必要があります。学習可能なのは `self.transformer`（デュアルストリーム Wan transformer）のみです。上記の例では、その自己注意（`attn1`）と相互注意（`attn2`）ブロックの両方のクエリ/バリュー射影を適応させており、これは標準的な LoRA ターゲットセットです。より高容量のアダプタが必要な場合は、範囲を広げてください（例：`to_k`/`to_out`、または `ffn` レイヤーを追加する）。`--peft.method_type` を渡すと PEFT が有効になるため、`--policy.use_peft=true` は不要です。

データセットは、カメラクリップ（各カメラの時間ウィンドウ、`frame_chunk_size`個の潜在フレームにVAEエンコード）と、各アイテムあたり`frame_chunk_size * action_per_frame`のアクションステップを提供する必要があります。

## データ形式（アクションチャネルとカメラの順序）

LingBot-VA は**エンドエフェクタ（デカルト）姿勢**のポリシーであり、関節位置ではなく EEF 姿勢とグリップを予測します。アクションは固定のマルチエンボディメント**30 次元**レイアウトに存在します。ロボットのアクション次元をこれらのチャネルにマッピングし、残りは`0`でパディングします（`used_action_channel_ids`は、特定のチェックポイントが実際に使用するチャネルを選択します）：

| チャンネル | 意味                                               |
| -------- | ----------------------------------------------------- |
| 0–6      | 左腕の末端エフェクタの姿勢                            |
| 7–13     | 右腕の末端エフェクタの姿勢                           |
| 14–20    | 左腕の関節（公開されたチェックポイントでは未使用）  |
| 21–27    | 右腕の関節（公開されたチェックポイントでは未使用） |
| 28       | 左グリップ                                          |
| 29       | 右グリップ                                         |

- **LIBERO** はチャンネル `0–6` を使用します: 6-DoF EEF デルタ (xyz + 回転) + グリッパー (単一アーム)。
- **RoboTwin**はチャンネル`[0–6, 28, 7–13, 29]`を使用します。左EEF（xyz + クォータニオン）+ 左グリップ + 右EEF + 右グリップ（16次元）です。環境はこれらのポーズをCuRobo IKを介して関節軌跡に変換します。関節は予測されません。

Joint-space datasets（または別のEEF規約）は、これらのチェックポイントのファインチューニングの前に、このスキーマに再マッピングする必要があります。

**カメラの順序は固定されており、順序に依存します**。各カメラの潜在変数は`obs_cam_keys`の順序で空間的に連結されるため、物理的なカメラとスロットの対応関係は学習時と一致させる必要があります。

| ベンチマーク | `obs_cam_keys`（順不同）                                                                             | `camera_layout`                                                     |
| --------- | ----------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------- |
| LIBERO    | `observation.images.image`（エージェント視点 / 第三者視点）、`observation.images.image2`（ハンドマウント / 手首）  | `width_concat`（幅方向に連結された latents）                      |
| RoboTwin  | `observation.images.head_camera`, `observation.images.left_camera`, `observation.images.right_camera` | `robotwin_tshape`（下部にフル解像度のヘッド、上部に2つの半解像度のリスト） |

最初のカメラは外観/ヘッドビューであり、残りはリストビューです。

## 推論ハイパーパラメータ (LIBERO)

| キー                                    | 値                                                                             |
| -------------------------------------- | --------------------------------------------------------------------------------- |
| 高さ × 幅                         | 128 × 128                                                                         |
| カメラ                                | `observation.images.image`（agentview）、`observation.images.image2`（eye-in-hand） |
| 使用されたアクションチャネル                   | 0–6（7自由度アーム＋グリップ）                                                         |
| action_per_frame / frame_chunk_size    | 4 / 4                                                                             |
| attn_window                            | 30                                                                                |
| video / action 去噪ステップ         | 20 / 50                                                                           |
| guidance_scale / action_guidance_scale | 5 / 1                                                                             |
| snr_shift / action_snr_shift           | 5.0 / 0.05                                                                        |

これらは`LingBotVAConfig`のデフォルト値です。`--policy.<name>=...`を介して、これらを上書きできます。

## 注記

- **Attention バックエンド:** 推論には `torch` SDPA バックエンド（常に利用可能）が使用されます。`flashattn` および `flex` バックエンドはオプションです。`flex` は学習にのみ必要です。
- **モデルサイズ:** DiTは約50億パラメータ、凍結されたVAE+UMT5は約20 GBを追加します。推論には、おおよそ18〜24 GBのVRAMが必要です。

## ライセンス

LingBot-VAはApache-2.0の下で公開されています。[アップストリームリポジトリ](https://github.com/Robbyant/lingbot-va)を参照してください。

