pypose/bae

📦 `bae` – ロボティクス向け2次スパース最適化

何であるか – bae（Bundle Adjustment Engine）は、PyTorchベースのライブラリであり、ロボティクスおよびビジョンで現れる大規模スパース問題、特にバンドル調整（BA） および ポーズグラフ最適化（PGO） に対して、正確な 2次最適化ソルバー（Levenberg-Marquardt、Schur補完）を実装しています。このライブラリはPyPoseエコシステムに基づき、そのLie群型とスパースヤコビアンユーティリティを再利用し、高性能スパース線形代数用のカスタムCUDAカーネルを追加しています。

🎯 コア機能

機能	重要性
スパースブロック行列演算	複数のカメラが3次元点を共有する際に生じるブロック構造ヤコビアンを処理し、メモリ使用量を低く抑えます。
CUDA加速	カスタムカーネルにより、BA/PGO実行時間の大部分を占める線形方程式の解法（PCG、Schur補完）を高速化します。
PyTorch統合	モデルは通常の `torch.nn.Module` であり、自動微分はライブラリの `sjac=True` パラメータと併用可能です。
Levenberg-Marquardt最適化	伝統的な非線形最小二乗ソルバーであり、今や完全なGPUサポートを備えています。
Schur補完ソルバー	線形システムをカメラ変数のみに簡略化し、非常に大きなシーンにおいてメモリ使用量を劇的に削減します。
`torch.compile`サポート	シムにより、BA残差とスパースヤコビアンの走査をTorchDynamo/Inductorでコンパイル可能にし、ミリ秒未満のカーネル起動を実現します。
エージェントスキルファイル	`.agent/skills/bae-compute-graph` は、LLM駆動エージェントが `bae` の上にカスタム計算グラフを構築するためのレシピを提供します。

🚀 一般的な利用事例

構造から運動（SfM）パイプライン – フィードフォワード予測器（例：Depth Anything 3、VGGT）の後にカメラポーズと3次元点を最適化。
SLAMバックエンド – 大きなループと高ノイズを伴うポーズグラフ問題を解決。
ロボティクスの認識 – 高密度RGB-D点群の整合（RGB-Dポーズ精緻化）をリアルタイムマッピングに統合。
研究プロトタイピング – 新しい残差や多様体パラメータ化を実験しながら、正確な2次更新を維持。

📦 インストール（クイックスタート）

# GitHubから直接ライブラリをインストール
python -m pip install git+https://github.com/pypose/bae.git

前提条件

CUDA 12.x（または13.x）ツールキット
同じCUDAバージョンでコンパイルされたPyTorch 2.0+
オプション：スパースソルバー加速用のNVIDIA CUDSS（pip install "nvidia-cudss-cu12<0.9"）
PyPose（スパースバックエンド）– pip install git+https://github.com/pypose/pypose.git

ソースを編集したい開発者向けには、READMEに完全な開発環境セットアップ手順（クローン、pip install -r requirements.txt、編集可能モードでインストールなど）が記載されています。

🛠️ 最小限のバンドル調整例

import torch, pypose as pp
from pypose.autograd.function import psjac
from bae.optim import LM
from bae.utils.pysolvers import PCG
from datapipes.bal_loader import get_problem

class Reproj(torch.nn.Module):
    def __init__(self, cam, pts):
        super().__init__()
        self.pose   = pp.Parameter(cam, sjac=True)   # カメラブロック
        self.points = pp.Parameter(pts, sjac=True)   # 点ブロック
        self.pose.trim_SE3_grad = True

    @psjac                                 # 構造化ヤコビアンを有効化
    def project(points, cam):
        proj = pp.SE3(cam[..., :7]).Act(points)
        proj = -proj[..., :2] / proj[..., 2:3]
        f, k1, k2 = cam[..., -3], cam[..., -2], cam[..., -1]
        n = (proj**2).sum(-1, keepdim=True)
        r = 1 + k1 * n + k2 * n**2
        return proj * r * f

    def forward(self, obs, cidx, pidx):
        pred = Reproj.project(self.points[pidx], self.pose[cidx])
        return pred - obs

# BAL問題をロード（例：problem-49-7776-pre）
prob = get_problem("problem-49-7776-pre", "ladybug", use_quat=True)
prob = {k: v.cuda() for k, v in prob.items() if isinstance(v, torch.Tensor)}

model = Reproj(prob['camera_params'].clone(), prob['points_3d'].clone()).cuda()

strategy = pp.optim.strategy.TrustRegion(up=2.0, down=0.5**4)
solver   = PCG(tol=1e-4, maxiter=250)
optim    = LM(model, strategy=strategy, solver=solver, reject=30)

for it in range(20):
    loss = optim.step({
        "observes": prob['points_2d'],
        "cidx": prob['camera_index_of_observations'],
        "pidx": prob['point_index_of_observations'],
    })
    print(f"Iter {it:02d} – loss {loss.item():.6f}")

このループはGPU上で完全なLevenberg-Marquardtステップを実行し、カメラポーズ（SE(3) Lie群パラメータとして）と3次元点座標の両方を更新します。

🤖 エージェントスキルサポート

このリポジトリには、カスタム残差グラフ（例：IMU要因やループクロージャー制約の追加）を構築する方法を記述したスキル（.agent/skills/bae-compute-graph）が同梱されています。LLM駆動のコーディングアシスタントはこのスキルをインポートすることで、新しい最適化問題用のボイラープレートコードを自動生成できます。

📚 詳細なリソースと引用

論文 – Bundle Adjustment in the Eager Mode（IEEE T-RO 2026）。PDF: https://arxiv.org/abs/2409.12190
引用（BibTeX）はREADMEに学術利用用に提供されています。
プロジェクトページ – https://pypose.org/bae/（デモ、メモリ・パフォーマンス分析、VGGT統合へのリンクを含む）。

🏁 まとめ

bae は、現代のビジュアルSLAM、SfM、ロボティクス認識パイプラインの中心にあるスパースでブロック構造を持つ問題に対して、GPUネイティブで正確な2次最適化を提供します。PyTorchの自動微分とカスタムCUDAカーネルを活用することで、迅速なプロトタイピングとプロダクションレベルのパフォーマンスのギャップを埋めています。