IEEE Transactions on Robotics (T-RO), 2026

Zitong Zhan, Huan Xu, Zihang Fang, Xinpeng Wei, Yaoyu Hu, and Chen Wang

🌐 Project Page | 📄 PDF

Basic functionalities in bae have been supported by LM in PyPose as a sparse backend starting from v0.9.5. Please refer to this example and docs of psjac for shared APIs used by both libraries.

bae is a PyTorch-based library supporting exact 2nd-order optimization techniques. The library provides efficient implementations for sparse optimization problems in robotics, particularly Bundle Adjustment (BA) and Pose Graph Optimization (PGO).

Bundle Adjustment

Indoor	Bonsai	Kitchen

_{bae powering BA and global positioning in downstream system, InstantSfM.}

Pose Graph Optimization

Sphere Big Noise	Grid3D	Sphere (g2o)

RGB-D Pose Refinement

Monitor Office	Lecture Room	Kitchenette

_{bae refines camera poses predicted by Depth Anything 3 (DA3), bringing feed-forward 3D reconstructions into sharper alignment.}

News

• 2026-03-22: Added skills for coding agents to write custom compute graphs.
• 2025-12-12: Added a VGGT integration example.

Features

• Sparse Block Matrix Operations: Optimized implementations of sparse matrix operations for large-scale optimization
• CUDA Acceleration: Custom CUDA kernels for high-performance sparse linear algebra
• PyTorch Integration: Seamlessly integrates with PyTorch's automatic differentiation framework
• Levenberg-Marquardt Optimizer: The vanilla LM algorithm for non-linear least squares problems
• Schur Complement Optimizer: Solve a Schur-reduce linear system for optimized memory consumption

Future Plan

• Reduce runtime overhead using CUDA graph; compile fwd & backward with torch.compile() reducing latency from 10ms to 2.2ms (added in PR #41)
• Reduce kernel launching cost in PCG solver by kernel fusion
• Distributed Tensor (DTensor) and FSDP support for multi-GPU and distributed optimization (WIP in distributed_dt)
• Schur complement (added in PR #35)
• Add Apple Silicon GPU support, PyTorch PR WIP
• Community-requested features: Need support for a new framework, hardware backend, or more efficient memory allocation? Open an issue and tell us about your use case; we’d love to hear what you need.

Installation

Prerequisites

• CUDA toolkit (tested with CUDA 12.x)
• PyTorch (2.0+)
• (Optional) CUDSS (CUDA Sparse Solver library)

User Setup Instructions

python -m pip install git+https://github.com/pypose/bae.git

Developer Setup Instructions

1. (Optional) Install CUDSS with pip package manager.

   • cuDSS 0.8.x uses the new datatype API; older releases retain the legacy API at compile time.
   • For CUDA 12.x, install nvidia-cudss-cu12. We verified nvidia-cudss-cu12==0.6.0.5 and nvidia-cudss-cu12==0.7.1.6 work with bae:

   pip install "nvidia-cudss-cu12<0.9"

   • For CUDA 13.x, install nvidia-cudss-cu13<0.9:

   pip install "nvidia-cudss-cu13<0.9"

2. Install PyPose:

   pip install git+https://github.com/pypose/pypose.git

3. Clone this repository:

   git clone https://github.com/zitongzhan/bae.git
   cd bae

4. Install dependencies:

   pip install -r requirements.txt

5. Install the package in development mode:

   python -m pip install --no-build-isolation -v -e .  # following https://github.com/pytorch/pytorch

Agent Skills

This repo includes skills in .agent/skills:

• bae-compute-graph: guidance for defining BAL/PGO and more complex compute graphs

Example Usage

Bundle Adjustment

Bundle Adjustment optimizes camera poses and 3D point positions to minimize reprojection error. The following example shows how to perform BA using bae:

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


class Reproj(torch.nn.Module):
    def __init__(self, camera_params, points):
        super().__init__()
        self.pose = pp.Parameter(camera_params, sjac=True)
        self.points = pp.Parameter(points, sjac=True)
        self.pose.trim_SE3_grad = True
    
    # Define the projection residual with structured Jacobian support
    @psjac
    def project(points, camera_params):
        projection = pp.SE3(camera_params[..., :7]).Act(points)
        projection = -projection[..., :2] / projection[..., [2]]

        f = camera_params[..., [-3]]
        k1 = camera_params[..., [-2]]
        k2 = camera_params[..., [-1]]

        n = torch.sum(projection**2, axis=-1, keepdim=True)
        r = 1 + k1 * n + k2 * n**2
        return projection * r * f

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


# Load a problem from the BAL dataset
dataset = get_problem("problem-49-7776-pre", "ladybug", use_quat=True)
dataset = {k: v.to('cuda') for k, v in dataset.items() if isinstance(v, torch.Tensor)}

# Prepare input for the optimization
input = {
    "observes": dataset['points_2d'],
    "cidx": dataset['camera_index_of_observations'],
    "pidx": dataset['point_index_of_observations'],
}

# Initialize model with camera parameters and 3D points
model = Reproj(
    dataset['camera_params'].clone(),
    dataset['points_3d'].clone(),
).to('cuda')

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

# Run optimization for multiple iterations
for idx in range(20):
    loss = optimizer.step(input)
    print(f'Iteration {idx}, loss: {loss.item()}')

Integration with VGGT

bae is used as an optional Bundle Adjustment backend in our VGGT fork (Visual Geometry Grounded Transformer) to refine the camera poses, intrinsics, and 3D points predicted by VGGT before exporting a COLMAP reconstruction.

After installing bae, you can run VGGT's COLMAP export with BA enabled and bae selected as the solver:

python demo_colmap.py --scene_dir /path/to/scene --use_ba --implementation bae  # optional: --shared_camera

This command invokes prepare_bae(...) inside vggt/demo_colmap.py, which wraps VGGT tracks and predictions into bae.optim.LM and updates extrinsic, intrinsic, and points_3d in place before writing scene_dir/sparse/ in COLMAP format.

Citation

If you use bae in your research, please cite:

@article{zhan2026bundle,
  title = {Bundle Adjustment in the Eager Mode},
  author = {Zhan, Zitong and Xu, Huan and Fang, Zihang and Wei, Xinpeng and Hu, Yaoyu and Wang, Chen},
  journal = {IEEE Transactions on Robotics},
  year = {2026},
  url = {https://arxiv.org/abs/2409.12190}
}

Acknowledgements

The implementation draws inspiration from:

• PyPose for SE(3) pose representations
• GTSAM for reprojection jacobian concepts
• Ceres for manifold parameter update