Lie Spline Torch

This example shows how to fit splines to an existing discrete-time trajectory. Note that we also provide a function to do this (fit_uniform_splines_to_traj()), which is used in the next example. See the function for a more rigoruous solution to this problem.

import torch
import lietorch as lt
from lie_spline_torch.spline_common import compute_num_knots_needed
from lie_spline_torch.uniform_rd_bspline import UniformRdBSpline
from lie_spline_torch.uniform_so3_bspline import UniformSO3BSpline

# ... (Read in the trajectory data for some local frame as torch tensors) ...
# positions: [N, 3] positions of the local frame origin in the world frame
# rotations: [N, 4] rotations that take a point from the local frame to the world frame (quaternions in xyzw order)
# times:     [N] pose timestamps (in seconds)

# Set the spline parameters
k = 4 # The spline order (4 ensures C2 continuity)
delta_time = (times[1:] - times[:-1]).mean() * 2.0 # The knot time interval (set to twice the average period here)
time_start = times[0] # Time of the first knot (in seconds)
sparse = False # Whether to compute sparse grads (False here as all control points will be active at each fitting iter.)
use_cpp = False # Whether to use the CPU-based C++ backend (False here as GPU backend is faster if len(times) is large)
use_pypose = False # Whether to use PyPose for GPU backend and return a PyPose object (False as LieTorch is faster)

# Set the initial spline control points
# NOTE: To keep the example simple we're initalizing the control points to identity. In practice, the trajectory data
# can be used to obtain a better initialization.
num_knots = compute_num_knots_needed(time_start, times[-1], delta_time, k) # Num. knots needed to cover the timespan
pos_control_points = torch.zeros(num_knots, 3).double() # double dtype recommended for control points and eval times
rot_control_points = torch.zeros(num_knots, 4).double()
rot_control_points[:, 3] = 1.0 # Rotation control points are specified as quaternions in xyzw format

# Initialize the splines
position_spline = UniformRdBSpline(k, delta_time, time_start, pos_control_points, sparse)
rotation_spline = UniformSO3BSpline(k, delta_time, time_start, rot_control_points, sparse, use_cpp, use_pypose)

# Initialize the optimizer
optimizer = torch.optim.Adam(list(position_spline.parameters()) + list(rotation_spline.parameters()))

# Fit the splines to the trajectory
tolerance = 1e-4
max_iter = 1000
rotations = lt.SO3(rotations) # Convert the rotations to a LieTorch SO3 instance
for i in range(max_iter):
    optimizer.zero_grad()

    # Evaluate the splines
    eval_positions = position_spline(times)
    eval_rotations = rotation_spline(times) # Rotations are returned as a LieTorch SO3 instance

    # Compute a loss
    # NOTE: As the operations in the rotation loss calculation are between LieTorch objects the backward pass
    # through this loss produces tangent space gradients of the loss with respect to the evaluated rotations.
    loss_position = (positions - eval_positions).norm(dim=-1).mean()
    loss_rotation = (rotations.inv() * eval_rotations).log().norm(dim=-1).mean()
    loss = loss_position + loss_rotation

    if loss < tolerance:
        break
    loss.backward()
    optimizer.step()

Note that there are several intermediate values calculated in the evaluation of the splines that are independent of the control points. Therefore, these values can be used repeatedly in an optimization loop where the evaluation times remain constant. To exploit this, we provide the option to precompute these values and repeatedly use them for spline evaluation in the loop:

# ... (Initialize the splines and set `tolerance`, `max_iter`, and `rotations` the same as before) ...

# Precompute intermediate values that are independent of the control points
pos_intermediate_dict = position_spline.compute_all_intermediate_values(times)
rot_intermediate_dict = rotation_spline.compute_all_intermediate_values(times)

for i in range(max_iter):
    optimizer.zero_grad()

    # Evaluate the splines
    eval_positions = position_spline(None, pos_intermediate_dict)
    eval_rotations = rotation_spline(None, rot_intermediate_dict)

    # ... (Compute the loss and update the parameters the same as before) ...

Let's say you have discrete-time pose estimates from applying COLMAP to a video and you want to optimize the continuous-time camera trajectory jointly with a NeRF or 3DGS scene. You can do this as follows:

from lie_spline_torch.fit_splines_to_traj import fit_uniform_splines_to_traj

# ... (Read in the COLMAP data for the camera frame as numpy arrays) ...
# positions: [N, 3] positions of the camera frame origin in the world frame
# rotations: [N, 4] rotations that take a point from the camera frame to the world frame (quaternions in xyzw order)
# times:     [N] pose timestamps (in seconds)

# Initialize the splines by fitting them to the discrete pose estimates
# NOTE: `sparse=True` sets the returned splines to compute sparse gradients. This is essential to limit optimizer
# updates to only the active control points in each iteration. Additional spline parameters can optionally be specified
# through this function.
position_spline, rotation_spline = fit_uniform_splines_to_traj(positions, rotations, times, device='cuda', sparse=True)

# (Optional) Move the spline parameters to the CPU and set the SO(3) spline to use the CPU-based C++ backend. This is
# likely the most efficient option if a relatively small number of poses are evaluated in each training iteration. The
# spline parameters don't need to be on the CPU to use the C++ backend, but it avoids device transfers in the spline
# evaluation.
position_spline.to('cpu')
rotation_spline.to('cpu')
rotation_spline.use_cpp = True

# Initialize the spline optimizer
# NOTE: SparseAdam is used here to handle the sparse gradients.
spline_optimizer = torch.optim.SparseAdam(list(position_spline.parameters()) + list(rotation_spline.parameters()))

# Training loop
while training:
    spline_optimizer.zero_grad()

    # ... (Zero other gradients) ...

    # ... (Read in a batch of pixels or images with corresponding timestamps `render_times`) ...

    # Evaluate the poses at the timestamps of the images or pixels to render
    render_positions = position_spline(render_times)
    render_rotations = rotation_spline(render_times)

    # (If necessary) Combine the positions and rotations into transformation matrices that take a point from the camera
    # frame to the world frame
    # NOTE: In this case, the rotations were returned as a LieTorch SO3 instance and the `matrix()` method converts
    # them to 4x4 transformation matrices with the rotation matrices embedded and zero translation. Similarly, if the
    # SO(3) spline was constructed with `use_pypose=True` the rotations would be returned as a PyPose SO3Type LieTensor
    # instance and the `matrix()` method would convert them to 3x3 rotation matrices.
    transformations = render_rotations.matrix()
    transformations[:, :3, 3] = render_positions

    # ... (Use the poses in rendering and compute a loss) ...

    loss.backward()
    spline_optimizer.step()

    # ... (Step other optimizers) ...

This example shows the core steps to optimize a continuous-time trajectory in a SLAM system.

import math
import torch
import lietorch as lt
from lie_spline_torch.uniform_rd_bspline import UniformRdBSpline
from lie_spline_torch.uniform_so3_bspline import UniformSO3BSpline

# ... (Set the spline parameters `k`, `delta_time`, and `time_start`) ...

# Initialize the splines
# NOTE: the splines are initialized here with the minimal number of control points (equal to the spline order) all set
# to identity. This sets the initial pose at `time_start` to identity.
pos_control_points = torch.zeros(k, 3).double()
rot_control_points = torch.zeros(k, 4).double()
rot_control_points[:, 3] = 1.0
position_spline = UniformRdBSpline(k, delta_time, time_start, pos_control_points, sparse=True)
rotation_spline = UniformSO3BSpline(k, delta_time, time_start, rot_control_points,
    sparse=True, use_cpp=True, use_pypose=False)

# Initialize the spline optimizer
spline_optimizer = torch.optim.SparseAdam(list(position_spline.parameters()) + list(rotation_spline.parameters()))

# Run SLAM
time_previous = time_start
while True:
    # ... (Read in new sensor data spanning `times_recent`) ...

    ##################################################    TRACKING    ##################################################

    # Set `time_current` to the latest timestamp in the new sensor data
    time_current = times_recent.max().item()

    ## Extend the splines to cover `time_current` using one of three options

    # Option 1: Duplicate the most recent control point until the necessary number of knots is reached
    # NOTE: This option may be best initially as early velocity estimates might be poor.
    position_spline.extend_to_time(time_current, spline_optimizer) # Optimizer states are also extended
    rotation_spline.extend_to_time(time_current, spline_optimizer)

    # Option 2: Predict positions and rotations using a constant velocity model and fit the splines to them
    # NOTE: `fit_to_constant_velocity` has additional arguments including the option to specify an interval over which
    # to compute a finite difference velocity estimate (better if the instantaneous velocity estimates are noisy).
    position_spline.extend_to_time(time_current, spline_optimizer) # It's necessary to extend them first
    rotation_spline.extend_to_time(time_current, spline_optimizer)
    position_spline.fit_to_constant_velocity(time_previous, time_current)
    rotation_spline.fit_to_constant_velocity(time_previous, time_current)

    # Option 3: Predict positions and rotations with a custom method and fit the splines to them
    # NOTE: For simplicity, we use a constant velocity model here despite this being equivalent to Option 2. In
    # TRGS-SLAM we also make predictions by integrating IMU measurements.
    position_spline.extend_to_time(time_current, spline_optimizer) # It's necessary to extend them first
    rotation_spline.extend_to_time(time_current, spline_optimizer)

    num_predict = int(math.ceil((time_current - time_previous) / delta_time)) * 10
    times_predict = torch.linspace(time_previous, time_current, num_predict)
    time_previous = torch.tensor([time_previous])

    position_previous = position_spline(time_previous)
    linear_velocity_previous = position_spline.evaluate_velocities(time_previous) # Analytical time derivative
    predicted_positions = position_previous + (times_predict.unsqueeze(-1) - time_previous) * linear_velocity_previous
    position_spline.fit_control_points(times_predict, predicted_positions)

    rotation_previous = rotation_spline(time_previous)
    angular_velocity_previous = rotation_spline.evaluate_velocities(time_previous) # Analytical time derivative
    predicted_rotations = \
        rotation_previous * lt.SO3.exp((times_predict.unsqueeze(-1) - time_previous) * angular_velocity_previous)
    rotation_spline.fit_control_points(times_predict, predicted_rotations)

    ## Tracking optimization loop
    for i in range(max_num_tracking_iter):
        spline_optimizer.zero_grad()

        # ... (Zero other gradients) ...

        # Evaluate the poses at `times_recent`
        est_positions = position_spline(times_recent)
        est_rotations = rotation_spline(times_recent)

        # (If needed) Evaluate time derivatives of the splines
        # - UniformRdBSpline has methods `evaluate_velocities()` and `evaluate_accelerations()`
        # - UniformSO3BSpline has `evaluate_velocities()`  (higher order derivatives may be added in the future)

        # ... (Use the poses and derivatives in computing a loss) ...

        loss.backward()
        spline_optimizer.step()

        # ... (Step other optimizers) ...
        # NOTE: The map (e.g., a radiance field) may be kept frozen in tracking.

    time_previous = time_current

    ##################################################    MAPPING    ###################################################
    # ... (Perform mapping -- similar to the tracking optimization loop but involves older data and updates the map) ...