Certifiable Landmark SLAM on SE(3)

import numpy as np
import plotly.graph_objects as go

import gtsam
from gtsam.symbol_shorthand import L, R, T

np.set_printoptions(precision=4, suppress=True)
rng = np.random.default_rng(6)

size, spacing = 3, 2.0

# Boustrophedon ("lawnmower") visit order through the cube.
order, cell_index = [], {}
for z in range(size):
    for y in (range(size) if z % 2 == 0 else reversed(range(size))):
        forward = (y % 2 == 0) == (z % 2 == 0)
        for x in (range(size) if forward else reversed(range(size))):
            cell_index[(x, y, z)] = len(order)
            order.append((x, y, z))

positions = np.array(order, dtype=float) * spacing

# Yaw follows the direction of travel; roll and pitch vary so the orientations
# do not stay inside a single one-parameter subgroup.
truth = []
for k in range(len(order)):
    step = positions[min(k + 1, len(order) - 1)] - positions[max(k - 1, 0)]
    yaw = np.arctan2(step[1], step[0]) if np.linalg.norm(step[:2]) > 1e-9 else 0.0
    truth.append(
        gtsam.Pose3(
            gtsam.Rot3.RzRyRx(0.15 * np.sin(k), -0.12 * np.cos(k), yaw), positions[k]
        )
    )

# Landmarks sit off the cube, each seen from the poses within range.
landmarks = np.array([[-2.0, 1.0, 2.0], [6.0, 5.0, 1.0], [2.0, -3.0, 3.0]])
sight_range = 6.0

odometry = [(k, k + 1) for k in range(len(order) - 1)]
sightings = [
    (k, b)
    for k in range(len(order))
    for b in range(len(landmarks))
    if np.linalg.norm(landmarks[b] - positions[k]) < sight_range
]
print(f"{len(order)} poses, {len(landmarks)} landmarks")
print(f"{len(odometry)} odometry edges, {len(sightings)} sightings")

kappa, tau, nu = 100.0, 25.0, 40.0  # rotation, odometry, sighting precisions

# Isotropic Langevin on rotation has sigma = 1 / sqrt(2 kappa) in the
# small-dispersion limit; the translation blocks are plain Gaussians.
sigma_R = 1.0 / np.sqrt(2.0 * kappa)
sigma_t = 1.0 / np.sqrt(tau)
sigma_v = 1.0 / np.sqrt(nu)

graph = gtsam.NonlinearFactorGraph()
rot_noise = gtsam.noiseModel.Isotropic.Variance(3, 1.0 / kappa)

for i, j in odometry:
    relative = truth[i].between(truth[j])
    measured = gtsam.Pose3(
        relative.rotation().compose(gtsam.Rot3.Expmap(rng.normal(scale=sigma_R, size=3))),
        relative.translation() + rng.normal(scale=sigma_t, size=3),
    )
    graph.add(
        gtsam.FrobeniusBetweenFactorRot3(R(i), R(j), measured.rotation(), rot_noise)
    )
    graph.add(
        gtsam.RelativeTranslationFactor3(R(i), T(i), T(j), measured.translation(), tau)
    )

for k, b in sightings:
    # The sighting is expressed in the body frame of pose k.
    body = truth[k].transformTo(landmarks[b]) + rng.normal(scale=sigma_v, size=3)
    graph.add(gtsam.RelativeTranslationFactor3(R(k), T(k), L(b), body, nu))

print(f"rotation:  kappa = {kappa:g}, sigma = {np.rad2deg(sigma_R):.2f} deg")
print(f"odometry:  tau   = {tau:g}, sigma = {sigma_t:.2f} m")
print(f"sighting:  nu    = {nu:g}, sigma = {sigma_v:.2f} m")
print(f"{graph.size()} factors")

rank = 3  # p = d for the first staircase level

initial = gtsam.Values()
for index in range(len(truth)):
    perturbed = truth[index].rotation().compose(
        gtsam.Rot3.Expmap(rng.normal(scale=0.3, size=3))
    )
    initial.insert(R(index), perturbed.matrix().T)                   # d x p
    initial.insert(T(index), rng.normal(scale=0.5, size=(1, rank)))  # 1 x p

# A landmark is one unconstrained row, the same shape as a translation.
for b in range(len(landmarks)):
    initial.insert(L(b), rng.normal(scale=0.5, size=(1, rank)))

print("rotation slice: ", initial.atMatrix(R(0)).shape)
print("landmark slice: ", initial.atMatrix(L(0)).shape)

alm_params = gtsam.AugmentedLagrangianParams()
alm_params.maxIterations = 200
# The sightings make the rotation constraints harder to satisfy than in the
# pose-only case, so start the augmented Lagrangian at a stiffer penalty; the
# default leaves the lifted rotations about 1e-2 off orthogonal here.
alm_params.bclInitialPenalty = 300.0

params = gtsam.RiemannianStaircaseParams()
params.pMin = rank
params.pMax = 7
params.eta = 1e-3
params.setAlmParams(alm_params)

result = gtsam.RiemannianStaircaseOptimizer(graph, initial, params).optimize()

print(f"Certified:    {result.certified}")
print(f"Final rank:   {result.finalRank}")
print(f"Ranks tried:  {np.asarray(result.getRanksVisited()).astype(int)}")
print(f"Solver time:  {result.totalTime:.4f} s")

if not result.hasRoundedSolution():
    raise RuntimeError("The staircase did not return a rounded solution.")

rounded = result.roundedValues()
blocks = {i: rounded.atMatrix(R(i)).copy() for i in range(len(truth))}
rows = {i: rounded.atMatrix(T(i)).copy() for i in range(len(truth))}
marks = {b: rounded.atMatrix(L(b)).copy() for b in range(len(landmarks))}

# Certifiable methods optimize over O(d); undo a global reflection if present.
if sum(np.linalg.det(block) < 0.0 for block in blocks.values()) > len(truth) // 2:
    for block in blocks.values():
        block[:, -1] *= -1.0
    for row in list(rows.values()) + list(marks.values()):
        row[:, -1] *= -1.0

def project_to_so(matrix):
    """Return the closest proper rotation to a square matrix."""
    left, _, right_transpose = np.linalg.svd(matrix)
    correction = np.eye(matrix.shape[0])
    correction[-1, -1] = np.linalg.det(left @ right_transpose)
    return left @ correction @ right_transpose

def rotation_from_matrix(matrix):
    """Rebuild a Rot3 from a proper 3-by-3 rotation matrix."""
    return gtsam.Rot3(matrix)

estimate = gtsam.Values()
for index in range(len(truth)):
    estimate.insert(R(index), rotation_from_matrix(project_to_so(blocks[index].T)))
    estimate.insert(T(index), gtsam.Point3(rows[index][0]))
for b in range(len(landmarks)):
    estimate.insert(L(b), gtsam.Point3(marks[b][0]))

truth_values = gtsam.Values()
for index in range(len(truth)):
    truth_values.insert(R(index), truth[index].rotation())
    truth_values.insert(T(index), gtsam.Point3(truth[index].translation()))
for b in range(len(landmarks)):
    truth_values.insert(L(b), gtsam.Point3(landmarks[b]))

relaxation_bound = float(np.asarray(result.getCostPerLevel())[-1])
rounded_objective = graph.error(estimate)
gap = (rounded_objective - relaxation_bound) / max(relaxation_bound, 1e-12)

print(f"Relaxation lower bound: {relaxation_bound:.6f}")
print(f"Rounded objective:      {rounded_objective:.6f}")
print(f"Relative gap:           {gap:.3e}")
print(f"Objective at truth:     {graph.error(truth_values):.6f}")

# Each odometry edge contributes 2d residuals and each sighting d; the gauge
# removes one pose.
residuals = len(odometry) * 6 + len(sightings) * 3
parameters = len(truth) * 6 + len(landmarks) * 3 - 6
dof = residuals - parameters
print(f"residuals {residuals} - parameters {parameters} = {dof} dof")
print(f"Expected optimal cost (0.5 * chi^2): {0.5 * dof:.1f}")
print(f"Observed optimal cost:               {rounded_objective:.1f}")
print(f"Ratio:                               {rounded_objective / (0.5 * dof):.2f}")

def procrustes(source, target):
    """Best rigid alignment of source onto target."""
    source_centre, target_centre = source.mean(0), target.mean(0)
    A, B = source - source_centre, target - target_centre
    left, _, right_transpose = np.linalg.svd(A.T @ B)
    return (A @ (left @ right_transpose)) + target_centre

positions_est = np.array([estimate.atPoint3(T(i)) for i in range(len(truth))])
landmarks_est = np.array([estimate.atPoint3(L(b)) for b in range(len(landmarks))])

# The relaxation fixes no gauge, so align pose and landmark estimates together.
aligned = procrustes(np.vstack([positions_est, landmarks_est]),
                     np.vstack([positions, landmarks]))
aligned_poses = aligned[: len(positions)]
aligned_landmarks = aligned[len(positions) :]

figure = go.Figure()
figure.add_trace(go.Scatter3d(x=positions[:, 0], y=positions[:, 1], z=positions[:, 2],
                              mode="lines+markers", name="ground truth",
                              line=dict(color="#2a78d6", width=4), marker=dict(size=3)))
figure.add_trace(go.Scatter3d(x=aligned_poses[:, 0], y=aligned_poses[:, 1],
                              z=aligned_poses[:, 2], mode="lines+markers",
                              name="estimate", line=dict(color="#eb6834", width=4),
                              marker=dict(size=3)))
figure.add_trace(go.Scatter3d(x=landmarks[:, 0], y=landmarks[:, 1], z=landmarks[:, 2],
                              mode="markers", name="landmarks (truth)",
                              marker=dict(color="#2a78d6", size=7, symbol="square-open")))
figure.add_trace(go.Scatter3d(x=aligned_landmarks[:, 0], y=aligned_landmarks[:, 1],
                              z=aligned_landmarks[:, 2], mode="markers",
                              name="landmarks (estimate)",
                              marker=dict(color="#eb6834", size=7, symbol="x")))
figure.update_layout(width=720, height=600,
                     scene=dict(xaxis_title="x [m]", yaxis_title="y [m]",
                                zaxis_title="z [m]", aspectmode="data"))
figure.show()

error = np.linalg.norm(aligned_poses - positions, axis=1)
print(f"RMS position error after alignment: {np.sqrt((error ** 2).mean()):.4f} m")

ranks = np.asarray(result.getRanksVisited()).astype(int)
phases = {
    "QCQP build": np.asarray(result.getQcqpBuildTimePerLevel()),
    "local solve": np.asarray(result.getNlpTimePerLevel()),
    "verify": np.asarray(result.getVerifyTimePerLevel()),
}

timing = go.Figure()
for name, seconds in phases.items():
    timing.add_trace(go.Bar(x=[f"p = {level}" for level in ranks], y=seconds, name=name))
timing.update_layout(
    barmode="stack", title="Time per staircase level",
    xaxis_title="staircase rank", yaxis_title="seconds",
    height=380, margin=dict(l=0, r=0, t=40, b=0))
timing.show()

for name, seconds in phases.items():
    print(f"{name:>12}: {seconds.sum():.5f} s")