SE(2) pose graph

In this notebook, we solve a 2D pose graph optimization problem: estimating robot poses from noisy relative measurements.

Inputs: Relative pose measurements (odometry + loop closures) from g2o file
Outputs: Globally consistent robot trajectory

Features used:

• SE2Var for SE(2) robot poses

• @jaxls.Cost.factory with batched edge construction

• Real g2o dataset with ~3500 poses and loop closures

Show code cell source

Hide code cell source

import sys
from loguru import logger

logger.remove()
logger.add(sys.stdout, format="<level>{level: <8}</level> | {message}");

import pathlib

import jax
import jax.numpy as jnp
import jaxlie
import jaxls
import numpy as np

Loading the dataset

Parse the g2o file to extract poses and edges. Edges include both odometry (consecutive poses) and loop closures (revisited locations):

@jax.jit
def parse_precision_matrix(components: jax.Array) -> jax.Array:
    """Convert upper triangular components to sqrt precision matrix.

    Args:
        components: Upper triangular components of the precision matrix (6,)

    Returns:
        Upper Cholesky factor of the precision matrix (3, 3)
    """
    precision = jnp.zeros((3, 3))
    triu_indices = jnp.triu_indices(3)
    precision = precision.at[triu_indices].set(components)
    precision = precision + precision.T - jnp.diag(jnp.diag(precision))
    return jnp.linalg.cholesky(precision).T


def parse_g2o(path: pathlib.Path) -> dict:
    """Parse a 2D g2o file into poses and edges.

    Args:
        path: Path to the g2o file

    Returns:
        Dictionary with 'poses' (N, 3) array and 'edges' list of tuples
    """
    with open(path) as f:
        lines = f.readlines()

    poses = []  # (x, y, theta)
    edges = []  # (i, j, dx, dy, dtheta, precision_components)

    for line in lines:
        parts = line.strip().split()
        if not parts:
            continue

        if parts[0] == "VERTEX_SE2":
            _, idx, x, y, theta = parts
            poses.append((float(x), float(y), float(theta)))

        elif parts[0] == "EDGE_SE2":
            _, i, j = parts[:3]
            dx, dy, dtheta = map(float, parts[3:6])
            precision_comps = list(map(float, parts[6:]))
            edges.append((int(i), int(j), dx, dy, dtheta, np.array(precision_comps)))

    return {"poses": np.array(poses), "edges": edges}

# Load the Manhattan 3500 dataset.
g2o_path = pathlib.Path("./data/input_M3500_g2o.g2o")
data = parse_g2o(g2o_path)

num_poses = len(data["poses"])
num_edges = len(data["edges"])

# Count odometry vs loop closure edges.
odometry_edges = [(i, j) for i, j, *_ in data["edges"] if j == i + 1]
loop_closure_edges = [(i, j) for i, j, *_ in data["edges"] if j != i + 1]

print(f"Poses: {num_poses}")
print(
    f"Edges: {num_edges} ({len(odometry_edges)} odometry, {len(loop_closure_edges)} loop closures)"
)

Poses: 3500
Edges: 5453 (3499 odometry, 1954 loop closures)

Variables and costs

Use SE2Var for poses on SE(2). The between factor measures the relative pose between two nodes, weighted by a precision matrix.

The sqrt_precision matrix is the upper Cholesky factor of the information (precision) matrix \(\Lambda = \Sigma^{-1}\), where \(\Sigma\) is the measurement covariance. By multiplying the residual by this factor, we get a whitened residual whose squared norm equals the Mahalanobis distance: \(r^T \Lambda \, r = \|L^T r\|^2\).

# Create batched pose variables.
pose_vars = jaxls.SE2Var(id=jnp.arange(num_poses))


@jaxls.Cost.factory
def between_cost(
    vals: jaxls.VarValues,
    var_a: jaxls.SE2Var,
    var_b: jaxls.SE2Var,
    measured: jaxlie.SE2,
    sqrt_precision: jax.Array,
) -> jax.Array:
    """Cost for relative pose measurement between two poses."""
    T_a = vals[var_a]
    T_b = vals[var_b]
    # Error: measured^{-1} @ (T_a^{-1} @ T_b)
    residual = (measured.inverse() @ (T_a.inverse() @ T_b)).log()
    return sqrt_precision @ residual


@jaxls.Cost.factory(kind="constraint_eq_zero")
def anchor_cost(
    vals: jaxls.VarValues,
    var: jaxls.SE2Var,
    target: jaxlie.SE2,
) -> jax.Array:
    """Anchor the first pose to prevent gauge freedom."""
    return (vals[var].inverse() @ target).log()

# Build edge arrays for batched cost construction.
edge_i = jnp.array([e[0] for e in data["edges"]])
edge_j = jnp.array([e[1] for e in data["edges"]])

# Measured relative poses.
measured_poses = jaxlie.SE2.from_xy_theta(
    jnp.array([e[2] for e in data["edges"]]),
    jnp.array([e[3] for e in data["edges"]]),
    jnp.array([e[4] for e in data["edges"]]),
)

# Sqrt precision matrices.
precision_comps = jnp.array([e[5] for e in data["edges"]])
sqrt_precisions = jax.vmap(parse_precision_matrix)(precision_comps)

print(f"Batched edge arrays: {edge_i.shape[0]} edges")

Batched edge arrays: 5453 edges

Solving

Problem structure

The problem structure connects pose variables through between factors (odometry and loop closures). With ~3500 poses and ~5600 edges, the visualization automatically limits nodes for performance:

# Create costs using batched construction.
costs: list[jaxls.Cost] = [
    # All between factors in one batched call.
    between_cost(
        jaxls.SE2Var(id=edge_i),
        jaxls.SE2Var(id=edge_j),
        measured_poses,
        sqrt_precisions,
    ),
    # Anchor first pose.
    anchor_cost(
        jaxls.SE2Var(id=0),
        jaxlie.SE2.from_xy_theta(
            data["poses"][0, 0], data["poses"][0, 1], data["poses"][0, 2]
        ),
    ),
]

# Visualize the problem structure structure.
problem = jaxls.LeastSquaresProblem(costs, [pose_vars])
problem.show()

# Initial values from g2o file (odometry integration).
initial_poses = jaxlie.SE2.from_xy_theta(
    jnp.array(data["poses"][:, 0]),
    jnp.array(data["poses"][:, 1]),
    jnp.array(data["poses"][:, 2]),
)
initial_vals = jaxls.VarValues.make([pose_vars.with_value(initial_poses)])

# Analyze the problem.
problem = problem.analyze()

INFO     | Building optimization problem with 5454 terms and 3500 variables: 5453 costs, 1 eq_zero, 0 leq_zero, 0 geq_zero
INFO     | Vectorizing group with 5453 costs, 2 variables each: between_cost
INFO     | Vectorizing constraint group with 1 constraints (constraint_eq_zero), 1 variables each: augmented_anchor_cost

# Solve with Gauss-Newton (no trust region needed for this well-conditioned problem)
solution = problem.solve(initial_vals, trust_region=None)

INFO     | Augmented Lagrangian: initial snorm=0.0000e+00, csupn=0.0000e+00, max_rho=1.0000e+07, constraint_dim=3
INFO     |  step #0: cost=2634708.0000 lambd=0.0000 inexact_tol=1.0e-02
INFO     |      - between_cost(5453): 2634708.00000 (avg 161.05557)
INFO     |      - augmented_anchor_cost(1): 0.00000 (avg 0.00000)
INFO     |      accepted=True ATb_norm=3.32e+04 cost_prev=2634708.0000 cost_new=57103.6992
INFO     |  step #1: cost=57103.6797 lambd=0.0000 inexact_tol=1.0e-02
INFO     |      - between_cost(5453): 57103.67969 (avg 3.49066)
INFO     |      - augmented_anchor_cost(1): 0.01626 (avg 0.00542)
INFO     |      accepted=True ATb_norm=3.59e+03 cost_prev=57103.6992 cost_new=11494.6016
INFO     |  step #2: cost=11494.5996 lambd=0.0000 inexact_tol=1.0e-02
INFO     |      - between_cost(5453): 11494.59961 (avg 0.70265)
INFO     |      - augmented_anchor_cost(1): 0.00040 (avg 0.00013)
INFO     |      accepted=True ATb_norm=1.23e+03 cost_prev=11494.6016 cost_new=3330.8013
INFO     |  step #3: cost=3330.8010 lambd=0.0000 inexact_tol=1.0e-02
INFO     |      - between_cost(5453): 3330.80103 (avg 0.20361)
INFO     |      - augmented_anchor_cost(1): 0.00016 (avg 0.00005)
INFO     |      accepted=True ATb_norm=4.14e+02 cost_prev=3330.8013 cost_new=11190.6309
INFO     |  step #4: cost=11190.6309 lambd=0.0000 inexact_tol=1.0e-02
INFO     |      - between_cost(5453): 11190.63086 (avg 0.68407)
INFO     |      - augmented_anchor_cost(1): 0.00000 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.02e+03 cost_prev=11190.6309 cost_new=2311.1602
INFO     |  step #5: cost=2311.1599 lambd=0.0000 inexact_tol=1.0e-02
INFO     |      - between_cost(5453): 2311.15991 (avg 0.14128)
INFO     |      - augmented_anchor_cost(1): 0.00001 (avg 0.00000)
INFO     |      accepted=True ATb_norm=4.05e+02 cost_prev=2311.1602 cost_new=737.6276
INFO     |  step #6: cost=737.6276 lambd=0.0000 inexact_tol=1.0e-02
INFO     |      - between_cost(5453): 737.62756 (avg 0.04509)
INFO     |      - augmented_anchor_cost(1): 0.00002 (avg 0.00001)
INFO     |      accepted=True ATb_norm=1.60e+02 cost_prev=737.6276 cost_new=1370.3214
INFO     |  step #7: cost=1370.3215 lambd=0.0000 inexact_tol=1.0e-02
INFO     |      - between_cost(5453): 1370.32153 (avg 0.08377)
INFO     |      - augmented_anchor_cost(1): 0.00001 (avg 0.00000)
INFO     |      accepted=True ATb_norm=2.83e+02 cost_prev=1370.3214 cost_new=183.5986
INFO     |  step #8: cost=183.5986 lambd=0.0000 inexact_tol=1.0e-02
INFO     |      - between_cost(5453): 183.59865 (avg 0.01122)
INFO     |      - augmented_anchor_cost(1): 0.00001 (avg 0.00000)
INFO     |      accepted=True ATb_norm=2.77e+01 cost_prev=183.5986 cost_new=219.7089
INFO     |  step #9: cost=219.7089 lambd=0.0000 inexact_tol=8.6e-03
INFO     |      - between_cost(5453): 219.70892 (avg 0.01343)
INFO     |      - augmented_anchor_cost(1): 0.00002 (avg 0.00001)
INFO     |      accepted=True ATb_norm=8.32e+01 cost_prev=219.7089 cost_new=145.7681
INFO     |  step #10: cost=145.7681 lambd=0.0000 inexact_tol=8.6e-03
INFO     |      - between_cost(5453): 145.76811 (avg 0.00891)
INFO     |      - augmented_anchor_cost(1): 0.00002 (avg 0.00001)
INFO     |      accepted=True ATb_norm=1.88e+00 cost_prev=145.7681 cost_new=152.7706
INFO     |  step #11: cost=152.7706 lambd=0.0000 inexact_tol=4.6e-04
INFO     |      - between_cost(5453): 152.77055 (avg 0.00934)
INFO     |      - augmented_anchor_cost(1): 0.00000 (avg 0.00000)
INFO     |      accepted=True ATb_norm=2.68e+01 cost_prev=152.7706 cost_new=138.1770
INFO     |  step #12: cost=138.1770 lambd=0.0000 inexact_tol=4.6e-04
INFO     |      - between_cost(5453): 138.17703 (avg 0.00845)
INFO     |      - augmented_anchor_cost(1): 0.00000 (avg 0.00000)
INFO     |      accepted=True ATb_norm=3.85e+00 cost_prev=138.1770 cost_new=137.9452
INFO     |  step #13: cost=137.9452 lambd=0.0000 inexact_tol=4.6e-04
INFO     |      - between_cost(5453): 137.94522 (avg 0.00843)
INFO     |      - augmented_anchor_cost(1): 0.00000 (avg 0.00000)
INFO     |      accepted=True ATb_norm=1.12e+00 cost_prev=137.9452 cost_new=137.9238
INFO     |  step #14: cost=137.9238 lambd=0.0000 inexact_tol=4.6e-04
INFO     |      - between_cost(5453): 137.92375 (avg 0.00843)
INFO     |      - augmented_anchor_cost(1): 0.00000 (avg 0.00000)
INFO     |      accepted=True ATb_norm=6.71e-01 cost_prev=137.9238 cost_new=137.9173
INFO     |  step #15: cost=137.9173 lambd=0.0000 inexact_tol=4.6e-04
INFO     |      - between_cost(5453): 137.91733 (avg 0.00843)
INFO     |      - augmented_anchor_cost(1): 0.00000 (avg 0.00000)
INFO     |      accepted=True ATb_norm=6.42e-01 cost_prev=137.9173 cost_new=137.9153
INFO     |  step #16: cost=137.9153 lambd=0.0000 inexact_tol=4.6e-04
INFO     |      - between_cost(5453): 137.91534 (avg 0.00843)
INFO     |      - augmented_anchor_cost(1): 0.00000 (avg 0.00000)
INFO     |      accepted=True ATb_norm=5.19e-01 cost_prev=137.9153 cost_new=137.9170
INFO     |  step #17: cost=137.9170 lambd=0.0000 inexact_tol=4.6e-04
INFO     |      - between_cost(5453): 137.91702 (avg 0.00843)
INFO     |      - augmented_anchor_cost(1): 0.00000 (avg 0.00000)
INFO     |      accepted=True ATb_norm=5.04e-01 cost_prev=137.9170 cost_new=137.9174
INFO     |  AL update: snorm=7.7177e-09, csupn=7.7177e-09, max_rho=1.0000e+07
INFO     | Terminated @ iteration #18: cost=137.9174 criteria=[1 0 0], term_deltas=3.1e-06,1.0e-01,6.6e-05 (solved in 0.7154 sec)

Visualization

Compare the initial odometry-only trajectory with the optimized result. Loop closures correct drift accumulated from odometry:

initial_xy = np.array(initial_vals[pose_vars].translation())
optimized_xy = np.array(solution[pose_vars].translation())

Show code cell source

Hide code cell source

import plotly.graph_objects as go
from plotly.subplots import make_subplots


def get_trajectory_trace(
    positions: np.ndarray,
    name: str,
    color: str,
) -> go.Scatter:
    """Create trajectory trace.

    Args:
        positions: Position array (N, 2)
        name: Trace name for legend
        color: Line/marker color

    Returns:
        Plotly Scatter trace
    """
    return go.Scatter(
        x=positions[:, 0],
        y=positions[:, 1],
        mode="lines",
        line=dict(color=color, width=1.5),
        name=name,
        hovertemplate="(%{x:.1f}, %{y:.1f})<extra></extra>",
    )


def get_loop_closure_traces(
    positions: np.ndarray,
    edges: list[tuple[int, int]],
    color: str,
    max_edges: int = 200,
) -> list[go.Scatter]:
    """Create loop closure edge traces (subsample for performance).

    Args:
        positions: Position array (N, 2)
        edges: List of (i, j) index pairs for loop closure edges
        color: Line color
        max_edges: Maximum number of edges to display

    Returns:
        List containing a single Plotly Scatter trace for all edges
    """
    # Subsample if too many.
    step = max(1, len(edges) // max_edges)
    sampled = edges[::step]

    x_coords = []
    y_coords = []
    for i, j in sampled:
        x_coords.extend([positions[i, 0], positions[j, 0], None])
        y_coords.extend([positions[i, 1], positions[j, 1], None])

    return [
        go.Scatter(
            x=x_coords,
            y=y_coords,
            mode="lines",
            line=dict(color=color, width=0.5),
            opacity=0.4,
            name="Loop closures",
            hoverinfo="skip",
        )
    ]

Show code cell source

Hide code cell source

from IPython.display import HTML

fig = make_subplots(
    rows=1,
    cols=2,
    subplot_titles=("Initial (odometry only)", "Optimized (with loop closures)"),
)

# Initial trajectory with loop closures shown.
for trace in get_loop_closure_traces(initial_xy, loop_closure_edges, "tomato"):
    fig.add_trace(trace, row=1, col=1)
fig.add_trace(get_trajectory_trace(initial_xy, "Trajectory", "steelblue"), row=1, col=1)

# Optimized trajectory.
for trace in get_loop_closure_traces(optimized_xy, loop_closure_edges, "tomato"):
    fig.add_trace(trace, row=1, col=2)
fig.add_trace(
    get_trajectory_trace(optimized_xy, "Trajectory", "forestgreen"), row=1, col=2
)

# Compute shared bounds.
all_xy = np.concatenate([initial_xy, optimized_xy])
x_min, x_max = all_xy[:, 0].min(), all_xy[:, 0].max()
y_min, y_max = all_xy[:, 1].min(), all_xy[:, 1].max()
padding = 0.05 * max(x_max - x_min, y_max - y_min)

# Layout with equal aspect and shared bounds.
fig.update_xaxes(
    title_text="x (m)",
    scaleanchor="y",
    scaleratio=1,
    range=[x_min - padding, x_max + padding],
)
fig.update_yaxes(
    title_text="y (m)",
    range=[y_min - padding, y_max + padding],
)
fig.update_layout(
    height=500,
    showlegend=False,
    margin=dict(t=40, b=40, l=40, r=40),
)
HTML(fig.to_html(full_html=False, include_plotlyjs="cdn"))

The optimization corrects drift accumulated from odometry-only integration. Loop closures (shown in red) constrain revisited locations to be consistent, resulting in a globally coherent map.

For more on Lie group variables, see jaxls.SE2Var and jaxls.SE3Var.