[WIP] Covariance Estimation

cpp/kiss_icp/core/Registration.cpp

@@ -29,6 +29,7 @@
 
 #include <algorithm>
 #include <cmath>
+#include <functional>
 #include <numeric>
 #include <sophus/se3.hpp>
 #include <sophus/so3.hpp>
@@ -41,8 +42,10 @@ using Matrix6d = Eigen::Matrix<double, 6, 6>;
 using Matrix3_6d = Eigen::Matrix<double, 3, 6>;
 using Vector6d = Eigen::Matrix<double, 6, 1>;
 }  // namespace Eigen
-using Associations = std::vector<std::pair<Eigen::Vector3d, Eigen::Vector3d>>;
-using LinearSystem = std::pair<Eigen::Matrix6d, Eigen::Vector6d>;
+
+using Correspondence = std::pair<Eigen::Vector3d, Eigen::Vector3d>;
+using Associations = std::vector<Correspondence>;
+using LinearSystem = std::tuple<Eigen::Matrix6d, Eigen::Vector6d, double>;
 
 namespace {
 inline double square(double x) { return x * x; }
@@ -120,7 +123,7 @@ Associations FindAssociations(const std::vector<Eigen::Vector3d> &points,
     return associations;
 }
 
-LinearSystem BuildLinearSystem(const Associations &associations, double kernel) {
+LinearSystem BuildLinearSystem(const Associations &associations, const double kernel_threshold) {
     auto compute_jacobian_and_residual = [](auto association) {
         const auto &[source, target] = association;
         const Eigen::Vector3d residual = source - target;
@@ -131,38 +134,59 @@ LinearSystem BuildLinearSystem(const Associations &associations, double kernel)
     };
 
     auto sum_linear_systems = [](LinearSystem a, const LinearSystem &b) {
-        a.first += b.first;
-        a.second += b.second;
+        std::get<Eigen::Matrix6d>(a) += std::get<Eigen::Matrix6d>(b);
+        std::get<Eigen::Vector6d>(a) += std::get<Eigen::Vector6d>(b);
+        std::get<double>(a) += std::get<double>(b);
         return a;
     };
 
-    auto GM_weight = [&](double residual2) { return square(kernel) / square(kernel + residual2); };
+    auto GM_kernel = [&](double residual2) {
+        return square(kernel_threshold) / square(kernel_threshold + residual2);
+    };
 
     using associations_iterator = Associations::const_iterator;
-    const auto &[JTJ, JTr] = tbb::parallel_reduce(
+    const auto &[JTJ, JTr, chi_square] = tbb::parallel_reduce(
         // Range
         tbb::blocked_range<associations_iterator>{associations.cbegin(), associations.cend()},
         // Identity
-        LinearSystem(Eigen::Matrix6d::Zero(), Eigen::Vector6d::Zero()),
+        LinearSystem(Eigen::Matrix6d::Zero(), Eigen::Vector6d::Zero(), 0.0),
         // 1st Lambda: Parallel computation
         [&](const tbb::blocked_range<associations_iterator> &r, LinearSystem J) -> LinearSystem {
             return std::transform_reduce(
                 r.begin(), r.end(), J, sum_linear_systems, [&](const auto &association) {
                     const auto &[J_r, residual] = compute_jacobian_and_residual(association);
-                    const double w = GM_weight(residual.squaredNorm());
-                    return LinearSystem(J_r.transpose() * w * J_r,        // JTJ
-                                        J_r.transpose() * w * residual);  // JTr
+                    const double chi_square = residual.squaredNorm();
+                    const double w = GM_kernel(chi_square);
+                    return LinearSystem(J_r.transpose() * w * J_r,                    // JTJ
+                                        J_r.transpose() * w * residual, chi_square);  // JTr
                 });
         },
         // 2nd Lambda: Parallel reduction of the private Jacboians
         sum_linear_systems);
 
-    return {JTJ, JTr};
+    return {JTJ, JTr, chi_square};
 }
 }  // namespace
 
 namespace kiss_icp {
 
+Estimate::Estimate() : pose(), covariance(Eigen::Matrix6d::Zero()) {}
+Estimate::Estimate(const Sophus::SE3d &T, const Eigen::Matrix6d &Sigma)
+    : pose(T), covariance(Sigma) {}
+
+// Equations from figure 3 of https://arxiv.org/pdf/1906.07795.pdf
+Estimate Estimate::inverse() const {
+    const auto Adj = pose.inverse().Adj();
+    return {pose.inverse(), Adj * covariance * Adj.transpose()};
+}
+
+Estimate operator*(Estimate lhs, const Estimate &rhs) {
+    lhs.pose *= rhs.pose;
+    const auto Adj = lhs.pose.Adj();
+    lhs.covariance += Adj * rhs.covariance * Adj.transpose();
+    return lhs;
+}
+
 Registration::Registration(int max_num_iteration, double convergence_criterion, int max_num_threads)
     : max_num_iterations_(max_num_iteration),
       convergence_criterion_(convergence_criterion),
@@ -174,35 +198,38 @@ Registration::Registration(int max_num_iteration, double convergence_criterion,
         tbb::global_control::max_allowed_parallelism, static_cast<size_t>(max_num_threads_));
 }
 
-Sophus::SE3d Registration::AlignPointsToMap(const std::vector<Eigen::Vector3d> &frame,
-                                            const VoxelHashMap &voxel_map,
-                                            const Sophus::SE3d &initial_guess,
-                                            double max_correspondence_distance,
-                                            double kernel) {
+Estimate Registration::AlignPointsToMap(const std::vector<Eigen::Vector3d> &frame,
+                                        const VoxelHashMap &voxel_map,
+                                        const Estimate &initial_guess,
+                                        double max_correspondence_distance,
+                                        double kernel_threshold) {
     if (voxel_map.Empty()) return initial_guess;
 
     // Equation (9)
     std::vector<Eigen::Vector3d> source = frame;
-    TransformPoints(initial_guess, source);
+    TransformPoints(initial_guess.pose, source);
 
     // ICP-loop
-    Sophus::SE3d T_icp = Sophus::SE3d();
+    Estimate icp_correction;
     for (int j = 0; j < max_num_iterations_; ++j) {
         // Equation (10)
-        const auto associations = FindAssociations(source, voxel_map, max_correspondence_distance);
+        const auto &associations = FindAssociations(source, voxel_map, max_correspondence_distance);
         // Equation (11)
-        const auto &[JTJ, JTr] = BuildLinearSystem(associations, kernel);
-        const Eigen::Vector6d dx = JTJ.ldlt().solve(-JTr);
+        const auto &[JTJ, JTr, chi_square] = BuildLinearSystem(associations, kernel_threshold);
+        const auto JTJ_inverse = JTJ.inverse();
+        const Eigen::Vector6d dx = -JTJ_inverse * JTr;
         const Sophus::SE3d estimation = Sophus::SE3d::exp(dx);
+        icp_correction.pose = estimation * icp_correction.pose;
+        // Follow Ollila et.al. https://arxiv.org/pdf/2006.10005.pdf - Eq. (2) for sample covariance
+        // using an M estimator
+        icp_correction.covariance = 1.0 / static_cast<double>(associations.size()) * JTJ_inverse;
         // Equation (12)
         TransformPoints(estimation, source);
-        // Update iterations
-        T_icp = estimation * T_icp;
         // Termination criteria
         if (dx.norm() < convergence_criterion_) break;
     }
     // Spit the final transformation
-    return T_icp * initial_guess;
+    return icp_correction * initial_guess;
 }
 
 }  // namespace kiss_icp

cpp/kiss_icp/core/Registration.hpp

@@ -28,16 +28,27 @@
 
 #include "VoxelHashMap.hpp"
 
+using CovarianceMatrixType = Eigen::Matrix<double, 6, 6>;
+
 namespace kiss_icp {
 
+struct Estimate {
+    Estimate();
+    Estimate(const Sophus::SE3d &T, const CovarianceMatrixType &Sigma);
+    Sophus::SE3d pose;
+    CovarianceMatrixType covariance;
+    Estimate inverse() const;
+    friend Estimate operator*(Estimate lhs, const Estimate &rhs);
+};
+
 struct Registration {
     explicit Registration(int max_num_iteration, double convergence_criterion, int max_num_threads);
 
-    Sophus::SE3d AlignPointsToMap(const std::vector<Eigen::Vector3d> &frame,
-                                  const VoxelHashMap &voxel_map,
-                                  const Sophus::SE3d &initial_guess,
-                                  double max_correspondence_distance,
-                                  double kernel);
+    Estimate AlignPointsToMap(const std::vector<Eigen::Vector3d> &frame,
+                              const VoxelHashMap &voxel_map,
+                              const Estimate &initial_guess,
+                              double max_correspondence_distance,
+                              double kernel_threshold);
 
     int max_num_iterations_;
     double convergence_criterion_;

cpp/kiss_icp/pipeline/KissICP.cpp

@@ -41,12 +41,12 @@ KissICP::Vector3dVectorTuple KissICP::RegisterFrame(const std::vector<Eigen::Vec
         // TODO(Nacho) Add some asserts here to sanitize the timestamps
 
         //  If not enough poses for the estimation, do not de-skew
-        const size_t N = poses().size();
+        const size_t N = estimates().size();
         if (N <= 2) return frame;
 
         // Estimate linear and angular velocities
-        const auto &start_pose = poses_[N - 2];
-        const auto &finish_pose = poses_[N - 1];
+        const auto &start_pose = estimates_[N - 2].pose;
+        const auto &finish_pose = estimates_[N - 1].pose;
         return DeSkewScan(frame, timestamps, start_pose, finish_pose);
     }();
     return RegisterFrame(deskew_frame);
@@ -64,19 +64,18 @@ KissICP::Vector3dVectorTuple KissICP::RegisterFrame(const std::vector<Eigen::Vec
 
     // Compute initial_guess for ICP
     const auto prediction = GetPredictionModel();
-    const auto last_pose = !poses_.empty() ? poses_.back() : Sophus::SE3d();
-    const auto initial_guess = last_pose * prediction;
-
+    const auto last_estimate = !estimates_.empty() ? estimates_.back() : Estimate();
+    const auto initial_guess = last_estimate * prediction;
     // Run icp
-    const Sophus::SE3d new_pose = registration_.AlignPointsToMap(source,         //
-                                                                 local_map_,     //
-                                                                 initial_guess,  //
-                                                                 3.0 * sigma,    //
-                                                                 sigma / 3.0);
-    const auto model_deviation = initial_guess.inverse() * new_pose;
+    const auto new_estimate = registration_.AlignPointsToMap(source,         //
+                                                             local_map_,     //
+                                                             initial_guess,  //
+                                                             3.0 * sigma,    //
+                                                             sigma / 3.0);
+    const auto model_deviation = (initial_guess.inverse() * new_estimate).pose;
     adaptive_threshold_.UpdateModelDeviation(model_deviation);
-    local_map_.Update(frame_downsample, new_pose);
-    poses_.push_back(new_pose);
+    local_map_.Update(frame_downsample, new_estimate.pose);
+    estimates_.push_back(new_estimate);
     return {frame, source};
 }
 
@@ -94,16 +93,19 @@ double KissICP::GetAdaptiveThreshold() {
     return adaptive_threshold_.ComputeThreshold();
 }
 
-Sophus::SE3d KissICP::GetPredictionModel() const {
-    Sophus::SE3d pred = Sophus::SE3d();
-    const size_t N = poses_.size();
-    if (N < 2) return pred;
-    return poses_[N - 2].inverse() * poses_[N - 1];
+Estimate KissICP::GetPredictionModel() const {
+    Estimate prediction;
+    const size_t N = estimates_.size();
+    if (N > 1) {
+        prediction.pose = (estimates_[N - 2].inverse() * estimates_[N - 1]).pose;
+    }
+    return prediction;
 }
 
 bool KissICP::HasMoved() {
-    if (poses_.empty()) return false;
-    const double motion = (poses_.front().inverse() * poses_.back()).translation().norm();
+    if (estimates_.empty()) return false;
+    const double motion =
+        (estimates_.front().pose.inverse() * estimates_.back().pose).translation().norm();
     return motion > 5.0 * config_.min_motion_th;
 }
 

cpp/kiss_icp/pipeline/KissICP.hpp

@@ -71,17 +71,17 @@ class KissICP {
                                       const std::vector<double> &timestamps);
     Vector3dVectorTuple Voxelize(const std::vector<Eigen::Vector3d> &frame) const;
     double GetAdaptiveThreshold();
-    Sophus::SE3d GetPredictionModel() const;
+    Estimate GetPredictionModel() const;
     bool HasMoved();
 
 public:
     // Extra C++ API to facilitate ROS debugging
     std::vector<Eigen::Vector3d> LocalMap() const { return local_map_.Pointcloud(); };
-    std::vector<Sophus::SE3d> poses() const { return poses_; };
+    std::vector<Estimate> estimates() const { return estimates_; };
 
 private:
     // KISS-ICP pipeline modules
-    std::vector<Sophus::SE3d> poses_;
+    std::vector<Estimate> estimates_;
     KISSConfig config_;
     Registration registration_;
     VoxelHashMap local_map_;

python/kiss_icp/kiss_icp.py

@@ -26,14 +26,14 @@
 from kiss_icp.deskew import get_motion_compensator
 from kiss_icp.mapping import get_voxel_hash_map
 from kiss_icp.preprocess import get_preprocessor
-from kiss_icp.registration import get_registration
+from kiss_icp.registration import Estimate, get_registration
 from kiss_icp.threshold import get_threshold_estimator
 from kiss_icp.voxelization import voxel_down_sample
 
 
 class KissICP:
     def __init__(self, config: KISSConfig):
-        self.poses = []
+        self.estimates = []
         self.config = config
         self.compensator = get_motion_compensator(config)
         self.adaptive_threshold = get_threshold_estimator(self.config)
@@ -43,7 +43,7 @@ def __init__(self, config: KISSConfig):
 
     def register_frame(self, frame, timestamps):
         # Apply motion compensation
-        frame = self.compensator.deskew_scan(frame, self.poses, timestamps)
+        frame = self.compensator.deskew_scan(frame, self.estimates, timestamps)
 
         # Preprocess the input cloud
         frame = self.preprocess(frame)
@@ -56,21 +56,22 @@ def register_frame(self, frame, timestamps):
 
         # Compute initial_guess for ICP
         prediction = self.get_prediction_model()
-        last_pose = self.poses[-1] if self.poses else np.eye(4)
-        initial_guess = last_pose @ prediction
+        initial_guess = self.estimates[-1] if self.estimates else Estimate()
+        initial_guess = initial_guess @ prediction
 
         # Run ICP
-        new_pose = self.registration.align_points_to_map(
+        new_estimate = self.registration.align_points_to_map(
             points=source,
             voxel_map=self.local_map,
             initial_guess=initial_guess,
             max_correspondance_distance=3 * sigma,
             kernel=sigma / 3,
         )
 
-        self.adaptive_threshold.update_model_deviation(np.linalg.inv(initial_guess) @ new_pose)
-        self.local_map.update(frame_downsample, new_pose)
-        self.poses.append(new_pose)
+        pose_deviation = (initial_guess.inverse() @ new_estimate).pose
+        self.adaptive_threshold.update_model_deviation(pose_deviation)
+        self.local_map.update(frame_downsample, new_estimate.pose)
+        self.estimates.append(new_estimate)
         return frame, source
 
     def voxelize(self, iframe):
@@ -86,13 +87,16 @@ def get_adaptive_threshold(self):
         )
 
     def get_prediction_model(self):
-        if len(self.poses) < 2:
-            return np.eye(4)
-        return np.linalg.inv(self.poses[-2]) @ self.poses[-1]
+        prediction = Estimate()
+        if len(self.estimates) > 1:
+            pose_from = self.estimates[-2].pose
+            pose_to = self.estimates[-1].pose
+            prediction.pose = np.linalg.inv(pose_from) @ pose_to
+        return prediction
 
     def has_moved(self):
-        if len(self.poses) < 1:
+        if len(self.estimates) < 1:
             return False
-        compute_motion = lambda T1, T2: np.linalg.norm((np.linalg.inv(T1) @ T2)[:3, -1])
-        motion = compute_motion(self.poses[0], self.poses[-1])
+        compute_motion = lambda T1, T2: np.linalg.norm((T1.inverse() @ T2).pose[:3, -1])
+        motion = compute_motion(self.estimates[0], self.estimates[-1])
         return motion > 5 * self.config.adaptive_threshold.min_motion_th

python/kiss_icp/pipeline.py

@@ -65,7 +65,9 @@ def __init__(
         self.odometry = KissICP(config=self.config)
         self.results = PipelineResults()
         self.times = []
-        self.poses = self.odometry.poses
+        self.estimates = self.odometry.estimates
+        self.poses = []
+        self.covariances = []
         self.has_gt = hasattr(self._dataset, "gt_poses")
         self.gt_poses = self._dataset.gt_poses[self._first : self._last] if self.has_gt else None
         self.dataset_name = self._dataset.__class__.__name__
@@ -89,6 +91,7 @@ def run(self):
         self._write_gt_poses()
         self._write_cfg()
         self._write_log()
+        self._write_covariances()
         return self.results
 
     # Private interface  ------
@@ -97,8 +100,12 @@ def _run_pipeline(self):
             raw_frame, timestamps = self._next(idx)
             start_time = time.perf_counter_ns()
             source, keypoints = self.odometry.register_frame(raw_frame, timestamps)
+            self.poses.append(self.estimates[-1].pose)
+            self.covariances.append(self.estimates[-1].covariance)
             self.times.append(time.perf_counter_ns() - start_time)
-            self.visualizer.update(source, keypoints, self.odometry.local_map, self.poses[-1])
+            self.visualizer.update(
+                source, keypoints, self.odometry.local_map, self.estimates[-1].pose
+            )
 
     def _next(self, idx):
         """TODO: re-arrange this logic"""
@@ -165,6 +172,9 @@ def _write_gt_poses(self):
             timestamps=self._get_frames_timestamps(),
         )
 
+    def _write_covariances(self):
+        np.save(f"{self.results_dir}/covariances.npy", np.stack(self.covariances))
+
     def _run_evaluation(self):
         # Run estimation metrics evaluation, only when GT data was provided
         if self.has_gt: