# bae **Repository Path**: gotoeasy/bae ## Basic Information - **Project Name**: bae - **Description**: No description available - **Primary Language**: Unknown - **License**: Apache-2.0 - **Default Branch**: release - **Homepage**: None - **GVP Project**: No ## Statistics - **Stars**: 0 - **Forks**: 0 - **Created**: 2026-06-12 - **Last Updated**: 2026-06-12 ## Categories & Tags **Categories**: Uncategorized **Tags**: None ## README bundle adjustment in the eager-mode

Accepted to IEEE Transactions on Robotics (T-RO), 2026

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

🌐 Project Page | 📄 PDF

> **⚠️ User Notice**: `bae` has been supported by [LM](https://pypose.org/docs/main/generated/pypose.optim.LevenbergMarquardt/#pypose.optim.LevenbergMarquardt) in [PyPose](https://github.com/pypose/pypose) as a sparse backend and is available from [v0.9.5](https://pypi.org/project/pypose/) or higher. Please refer to [this example](https://github.com/pypose/pypose/tree/main/examples/module/ba) and docs of [psjac](https://pypose.org/docs/main/generated/pypose.autograd.function.parallel_for_sparse_jacobian/#pypose.autograd.function.parallel_for_sparse_jacobian), a shared API for 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

Garden bundle adjustment example

Counter bundle adjustment example

Kitchen bundle adjustment example

Garden Counter Kitchen

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

### Pose Graph Optimization
Sphere big-noise optimization 3D grid optimization Sphere g2o optimization
Sphere Big Noise Grid3D Sphere (g2o)
## News - 2026-03-22: Added [skills](.agent/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 - **Bundle Adjustment**: Efficient implementation for camera pose and 3D structure optimization - **Pose Graph Optimization**: Tools for optimizing robot trajectories using pose graph representations - **PyTorch Integration**: Seamlessly integrates with PyTorch's automatic differentiation framework - **Levenberg-Marquardt Optimizer**: Custom implementation of the LM algorithm for non-linear least squares problems ### Future Plan - [ ] Add Apple Silicon GPU support, [PyTorch PR WIP](https://github.com/pytorch/pytorch/pull/177757) - [ ] Reduce runtime overhead using CUDA graph - [ ] Distributed Tensor (DTensor) and FSDP support for multi-GPU and distributed optimization - [ ] An new backend for [distributed solver](https://github.com/NVIDIA/AMGX) - [x] Schur complement (added in [PR #35](https://github.com/pypose/bae/pull/35)) ## Installation ### Prerequisites - CUDA toolkit (tested with CUDA 12.x) - PyTorch (2.0+) - (Optional) [CUDSS](https://developer.nvidia.com/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. - 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`: ```bash pip install "nvidia-cudss-cu12<=0.7.1.6" ``` - For CUDA 13.x, install `nvidia-cudss-cu13<=0.7.1.6`: ```bash pip install "nvidia-cudss-cu13<=0.7.1.6" ``` 2. Install PyPose: ```bash pip install git+https://github.com/pypose/pypose.git ``` 3. Clone this repository: ```bash git clone https://github.com/zitongzhan/bae.git cd bae ``` 4. Install dependencies: ```bash pip install -r requirements.txt ``` 5. Install the package in development mode: ```bash python -m pip install --no-build-isolation -v -e . # following https://github.com/pytorch/pytorch ``` ## Agent Skills This repo includes skills in [.agent/skills](.agent/skills): - [`bae-compute-graph`](.agent/skills/bae-compute-graph/SKILL.md): 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`: ```python 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](https://github.com/zitongzhan/vggt) (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: ```bash 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: ```bibtex @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](https://github.com/pypose/pypose) for SE(3) pose representations - GTSAM for reprojection jacobian concepts - Ceres for manifold parameter update