pypose/bae

📦 `bae` – 로보틱스를 위한 2차 스파스 최적화

무엇인가요 – bae (Bundle Adjustment Engine)는 로보틱스 및 비전에서 나타나는 대규모 스파스 문제, 특히 버블 조정(BA) 와 포즈 그래프 최적화(PGO) 에 대해 정확한 2차 최적화 솔버(Levenberg-Marquardt, Schur 보완)를 구현하는 PyTorch 기반 라이브러리입니다. 이 라이브러리는 PyPose 생태계를 기반으로 하며, Lie-군 타입과 스파스 자코비안 유틸리티를 재사용하고, 고성능 스파스 선형 대수를 위한 커스텀 CUDA 커널을 추가합니다.

🎯 핵심 기능

기능	왜 중요한가
스파스 블록 행렬 연산	여러 카메라가 3D 점을 공유할 때 나타나는 블록 구조 자코비안을 처리하여 메모리 사용량을 낮춥니다.
CUDA 가속	커스텀 커널이 선형 해법(PCG, Schur 보완)을 가속화하여 BA/PGO 실행 시간의 대부분을 차지합니다.
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) 이후 카메라 자세와 3D 점을 정밀 조정.
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-군 파라미터로)와 3D 점 좌표를 모두 업데이트합니다.

🤖 에이전트 스킬 지원

이 리포지토리는 사용자 정의 잔차 그래프(예: 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 커널을 활용하여, 빠른 프로토타이핑과 프로덕션 수준 성능 사이의 격차를 메웁니다.