pypose/bae

📦 `bae` – 2nd‑order Sparse Optimization for Robotics

What it is – bae (Bundle Adjustment Engine) is a PyTorch‑based library that implements exact second‑order solvers (Levenberg‑Marquardt, Schur‑complement) for large sparse problems that appear in robotics and vision, most notably Bundle Adjustment (BA) and Pose‑Graph Optimization (PGO). It builds on the PyPose ecosystem, re‑using its Lie‑group types and sparse‑Jacobian utilities, and adds custom CUDA kernels for high‑performance sparse linear algebra.

🎯 Core capabilities

Feature	Why it matters
Sparse block‑matrix ops	Handles the block‑structured Jacobians that arise when many cameras share 3‑D points, keeping memory usage low.
CUDA acceleration	Custom kernels speed up the linear solves (PCG, Schur complement) that dominate BA/PGO runtimes.
PyTorch integration	Models are ordinary `torch.nn.Module`s; automatic differentiation works with the library’s `sjac=True` parameters.
Levenberg‑Marquardt optimizer	Classic non‑linear least‑squares solver, now available with full GPU support.
Schur‑complement solver	Reduces the linear system to camera variables only, dramatically cutting memory for very large scenes.
`torch.compile` support	A shim lets you compile BA residuals and sparse‑Jacobian traversals with TorchDynamo/Inductor, yielding sub‑millisecond kernel launches.
Agent‑skill files	`.agent/skills/bae‑compute‑graph` gives LLM‑driven agents a recipe for building custom compute graphs on top of `bae`.

🚀 Typical use‑cases

Structure‑from‑Motion pipelines – refine camera poses and 3‑D points after a feed‑forward predictor (e.g., Depth Anything 3, VGGT).
SLAM back‑ends – solve pose‑graph problems with large loops and high noise.
Robotics perception – integrate dense RGB‑D point‑cloud alignment (RGB‑D pose refinement) into real‑time mapping.
Research prototyping – experiment with new residuals or manifold parameterizations while still getting exact second‑order updates.

📦 Installation (quick start)

# Install the library directly from GitHub
python -m pip install git+https://github.com/pypose/bae.git

Prerequisites

CUDA 12.x (or 13.x) toolkit
PyTorch 2.0+ compiled with the same CUDA version
Optional: NVIDIA CUDSS for sparse‑solver acceleration (pip install "nvidia-cudss-cu12<0.9")
PyPose (the sparse‑backend) – pip install git+https://github.com/pypose/pypose.git

For developers who want to edit the source, the README provides a full dev‑setup (clone, pip install -r requirements.txt, install in editable mode, etc.).

🛠️ Minimal Bundle‑Adjustment example

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)   # camera blocks
        self.points = pp.Parameter(pts, sjac=True)   # point blocks
        self.pose.trim_SE3_grad = True

    @psjac                                 # enables structured Jacobian
    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

# Load a BAL problem (e.g., 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}")

The loop runs a full Levenberg‑Marquardt step on the GPU, updating both camera poses (as SE(3) Lie‑group parameters) and 3‑D point coordinates.

🤖 Agent‑skill support

The repository ships a skill (.agent/skills/bae-compute-graph) that describes how to construct custom residual graphs (e.g., adding IMU factors or loop‑closure constraints). LLM‑driven coding assistants can import this skill to automatically generate boiler‑plate code for new optimization problems.

📚 Further reading & citation

Paper – Bundle Adjustment in the Eager Mode (IEEE T‑RO 2026). PDF: https://arxiv.org/abs/2409.12190
Citation (BibTeX) is provided in the README for academic use.
Project page – https://pypose.org/bae/ (contains demos, memory‑performance analysis, and links to the VGGT integration).

🏁 Bottom line

bae gives researchers and engineers a GPU‑native, exact second‑order optimizer for the sparse, block‑structured problems that are at the heart of modern visual‑SLAM, SfM, and robotics perception pipelines. By leveraging PyTorch’s autograd and custom CUDA kernels, it bridges the gap between rapid prototyping and production‑grade performance.