pypose/bae

📦 `bae` – 机器人学中的二阶稀疏优化

是什么 – bae（Bundle Adjustment Engine）是一个基于 PyTorch 的库，实现了针对机器人学和视觉中出现的大规模稀疏问题的精确二阶求解器（Levenberg-Marquardt、Schur 补）—— 最著名的是 捆绑调整（BA） 和 位姿图优化（PGO）。它基于 PyPose 生态系统，复用其李群类型和稀疏雅可比工具，并添加了自定义 CUDA 内核以实现高性能稀疏线性代数。

🎯 核心功能

特性	为何重要
稀疏块矩阵运算	处理由多个相机共享 3D 点时产生的块结构雅可比矩阵，保持低内存占用。
CUDA 加速	自定义内核加速了主导 BA/PGO 运行时间的线性求解（PCG、Schur 补）。
PyTorch 集成	模型是普通的 `torch.nn.Module`；自动微分可与库的 `sjac=True` 参数配合使用。
Levenberg-Marquardt 优化器	经典的非线性最小二乘求解器，现已支持完整 GPU 加速。
Schur 补求解器	将线性系统简化为仅相机变量，大幅降低超大场景的内存消耗。
`torch.compile` 支持	通过一个适配器，可使用 TorchDynamo/Inductor 编译 BA 残差和稀疏雅可比遍历，实现亚毫秒级内核启动。
Agent-技能文件	`.agent/skills/bae-compute-graph` 提供了 LLM 驱动的代理构建自定义计算图的配方。

🚀 典型应用场景

从运动中恢复结构（SfM）流水线 – 在前馈预测器（如 Depth Anything 3、VGGT）之后，优化相机位姿和 3D 点。
SLAM 后端 – 解决具有大回环和高噪声的位姿图问题。
机器人感知 – 将密集 RGB-D 点云对齐（RGB-D 位姿精炼）集成到实时建图中。
研究原型开发 – 在保持精确二阶更新的同时，实验新的残差或流形参数化。

📦 安装（快速入门）

# 直接从 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) 李群参数）和 3D 点坐标。

🤖 Agent-技能支持

该仓库附带一个技能（.agent/skills/bae-compute-graph），描述了如何构建自定义残差图（例如添加 IMU 因子或回环约束）。LLM 驱动的编码助手可导入此技能，自动生成新优化问题的样板代码。

📚 进一步阅读与引用

论文 – 在急切模式下的捆绑调整（IEEE T-RO 2026）。PDF: https://arxiv.org/abs/2409.12190
引用（BibTeX）在 README 中提供，供学术使用。
项目页面 – https://pypose.org/bae/（包含演示、内存-性能分析，以及与 VGGT 集成的链接）。

🏁 总结

bae 为研究人员和工程师提供了一种 GPU 原生、精确的二阶优化器，用于现代视觉 SLAM、SfM 和机器人感知流水线中的核心稀疏、块结构问题。通过利用 PyTorch 的自动微分和自定义 CUDA 内核，它弥合了快速原型开发与生产级性能之间的差距。