Repository: resim-ai/open-core
Pull request #240: Add two jet mean function
Author: login: mikebauer
Created: 2025-03-03T21:00:29Z | Merged: 2025-03-05T22:37:29Z

## Description of change
Two Jets exist on a manifold so their mean $\mu$ is computed iteratively such that:

$$
\sum_i (x_i \ominus \mu) = 0
$$

In our case $\ominus$ is defined by the difference in [two_jet_tangent_space](https://github.com/resim-ai/open-core/blob/main/resim/curves/optimization/two_jet_tangent_space.hh).

## Guide to reproduce test results.
```bash
bazel test //resim/curves/learning/...
```
## Checklist:
- [x] I have self-reviewed this change.
- [x] I have tested this change.
- [x] This change is covered by tests that are already landed, or in this PR.



Diff:
diff --git a/resim/curves/learning/BUILD b/resim/curves/learning/BUILD
new file mode 100644
index 00000000..040ef6ff
--- /dev/null
+++ b/resim/curves/learning/BUILD
@@ -0,0 +1,35 @@
+# Copyright 2025 ReSim, Inc.
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+
+load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test")
+
+cc_library(
+    name = "two_jet_mean",
+    srcs = ["two_jet_mean.cc"],
+    hdrs = ["two_jet_mean.hh"],
+    deps = [
+        "//resim/assert",
+        "//resim/curves:two_jet",
+        "//resim/curves/optimization:two_jet_tangent_space",
+        "//resim/transforms:se3",
+        "//resim/utils:status_value",
+    ],
+)
+
+cc_test(
+    name = "two_jet_mean_test",
+    srcs = ["two_jet_mean_test.cc"],
+    deps = [
+        ":two_jet_mean",
+        "//resim/curves:two_jet_test_helpers",
+        "//resim/curves/optimization:two_jet_tangent_space",
+        "//resim/testing:random_matrix",
+        "//resim/transforms:se3",
+        "//resim/transforms:so3",
+        "//resim/utils:inout",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
diff --git a/resim/curves/learning/two_jet_mean.cc b/resim/curves/learning/two_jet_mean.cc
new file mode 100644
index 00000000..8526f84e
--- /dev/null
+++ b/resim/curves/learning/two_jet_mean.cc
@@ -0,0 +1,47 @@
+// Copyright 2025 ReSim, Inc.
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file or at
+// https://opensource.org/licenses/MIT.
+
+#include "resim/curves/learning/two_jet_mean.hh"
+
+#include "resim/assert/assert.hh"
+#include "resim/curves/optimization/two_jet_tangent_space.hh"
+#include "resim/curves/two_jet.hh"
+
+namespace resim::curves::learning {
+
+using transforms::SE3;
+using TwoJetL = TwoJetL<SE3>;
+using TangentVector = optimization::TwoJetTangentVector<SE3>;
+
+StatusValue<TwoJetL> two_jet_mean(
+    const std::span<const TwoJetL> &samples,
+    const double tolerance,
+    const int max_iterations) {
+  REASSERT(not samples.empty(), "Sample size must be greater than 1");
+  REASSERT(
+      max_iterations > 0,
+      "Number of iterations must be greater than zero");
+
+  TwoJetL guess = samples.front();
+  TangentVector delta{TangentVector::Zero()};
+  bool converged = false;
+  for (int ii = 0; ii < max_iterations and not converged; ++ii) {
+    for (const auto &sample : samples) {
+      delta += optimization::difference(sample, guess);
+    }
+    delta /= static_cast<double>(samples.size());
+    guess = optimization::accumulate(guess, delta);
+    if (delta.norm() < tolerance) {
+      converged = true;
+    }
+  }
+  if (not converged) {
+    return MAKE_STATUS("Failed to converge!");
+  }
+  return guess;
+}
+
+}  // namespace resim::curves::learning
diff --git a/resim/curves/learning/two_jet_mean.hh b/resim/curves/learning/two_jet_mean.hh
new file mode 100644
index 00000000..5826d7ac
--- /dev/null
+++ b/resim/curves/learning/two_jet_mean.hh
@@ -0,0 +1,36 @@
+// Copyright 2025 ReSim, Inc.
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file or at
+// https://opensource.org/licenses/MIT.
+
+#pragma once
+
+#include <span>
+
+#include "resim/curves/two_jet.hh"
+#include "resim/transforms/se3.hh"
+#include "resim/utils/status_value.hh"
+
+namespace resim::curves::learning {
+
+// Compute the mean of a set of two jets. Since two jets exist on a manifold,
+// this function computes the two jet u such that sum_i(x_i (-) u) = 0 where (-)
+// is defined as the left minus operator as defined in
+// resim/curves/optimization/two_jet_tangent_space.hh. Another way of saying
+// this is that u is the mean of all the samples when they're expressed in u's
+// tangent space. We use iteration to find u and therefore a tolerance and a
+// maximum number of iterations must be provided.
+// @param[in] samples - The samples to compute the mean of.
+// @param[in] tolerance - The absolute tolerance for computing the mean.
+//                        Iteration terminates once updates get smaller in L2
+//                        norm than this.
+// @param[in] max_iterations - The maximum number of iterations allowed before
+//                             terminating unsuccessfully.
+// @returns The mean as described above, or a bad status otherwise.
+StatusValue<TwoJetL<resim::transforms::SE3>> two_jet_mean(
+    const std::span<const TwoJetL<resim::transforms::SE3>> &samples,
+    double tolerance,
+    int max_iterations);
+
+}  // namespace resim::curves::learning
diff --git a/resim/curves/learning/two_jet_mean_test.cc b/resim/curves/learning/two_jet_mean_test.cc
new file mode 100644
index 00000000..ae037ada
--- /dev/null
+++ b/resim/curves/learning/two_jet_mean_test.cc
@@ -0,0 +1,155 @@
+// Copyright 2025 ReSim, Inc.
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file or at
+// https://opensource.org/licenses/MIT.
+
+#include "resim/curves/learning/two_jet_mean.hh"
+
+#include <gtest/gtest.h>
+
+#include <algorithm>
+#include <random>
+
+#include "resim/curves/optimization/two_jet_tangent_space.hh"
+#include "resim/curves/two_jet_test_helpers.hh"
+#include "resim/testing/random_matrix.hh"
+#include "resim/transforms/se3.hh"
+#include "resim/utils/inout.hh"
+
+namespace resim::curves::learning {
+
+using transforms::SE3;
+using Vec3 = Eigen::Vector3d;
+using TwoJetL = TwoJetL<SE3>;
+using TangentVector = optimization::TwoJetTangentVector<SE3>;
+
+// Simple helper to create random sets of TwoJets drawn from a normal
+// distribution.
+template <typename Rng>
+std::vector<TwoJetL> random_samples(
+    const int num_samples,
+    const TwoJetL &dist_mean,
+    const double stddev,
+    InOut<Rng> rng) {
+  std::vector<TwoJetL> samples;
+  samples.reserve(num_samples);
+
+  std::normal_distribution<double> dist{0.0, stddev};
+  const auto generator = [&]() -> TwoJetL {
+    return optimization::accumulate(
+        dist_mean,
+        testing::random_matrix<TangentVector>(*rng, dist));
+  };
+
+  std::generate_n(std::back_inserter(samples), num_samples, generator);
+
+  return samples;
+}
+
+TEST(TwoJetMeanTest, TestTwoJetMean) {
+  // SETUP
+  constexpr std::size_t NUM_SAMPLES = 50;
+  constexpr int SEED = 8932;
+  constexpr double STDDEV = 0.1;
+  constexpr int MAX_ITERATIONS = NUM_SAMPLES;
+  constexpr int NUM_TESTS = 10;
+  constexpr double TOLERANCE = 1e-8;
+  constexpr double MEAN_TOLERANCE = 5e-2;
+
+  TwoJetTestHelper<TwoJetL> helper;
+  std::mt19937 rng{SEED};
+
+  for (int ii = 0; ii < NUM_TESTS; ++ii) {
+    auto dist_mean = helper.make_test_two_jet();
+    const auto samples =
+        random_samples(NUM_SAMPLES, dist_mean, STDDEV, InOut{rng});
+
+    // ACTION
+    const auto sample_mean_sv =
+        two_jet_mean(samples, TOLERANCE, MAX_ITERATIONS);
+
+    // VERIFICATION
+    ASSERT_TRUE(sample_mean_sv.ok());
+    const TwoJetL &sample_mean = sample_mean_sv.value();
+
+    EXPECT_TRUE(sample_mean.is_approx(dist_mean, MEAN_TOLERANCE));
+  }
+}
+
+TEST(TwoJetMeanTest, TestTwoJetMeanSingleSample) {
+  // SETUP
+  constexpr int SEED = 8932;
+  constexpr int MAX_ITERATIONS = 10;
+  constexpr int NUM_TESTS = 10;
+  constexpr double TOLERANCE = 1e-8;
+  constexpr double MEAN_TOLERANCE = 1e-12;
+
+  TwoJetTestHelper<TwoJetL> helper;
+  std::mt19937 rng{SEED};
+
+  for (int ii = 0; ii < NUM_TESTS; ++ii) {
+    auto dist_mean = helper.make_test_two_jet();
+    const std::vector<TwoJetL> samples = {dist_mean};
+
+    // ACTION
+    const auto sample_mean_sv =
+        two_jet_mean(samples, TOLERANCE, MAX_ITERATIONS);
+
+    // VERIFICATION
+    ASSERT_TRUE(sample_mean_sv.ok());
+    const TwoJetL &sample_mean = sample_mean_sv.value();
+
+    EXPECT_TRUE(sample_mean.is_approx(dist_mean, MEAN_TOLERANCE));
+  }
+}
+
+TEST(TwoJetMeanTest, TestTwoJetNoConverge) {
+  // SETUP
+  constexpr std::size_t NUM_SAMPLES = 50;
+  constexpr int SEED = 8932;
+  constexpr double STDDEV = 0.1;
+  constexpr int MAX_ITERATIONS = 1;
+  constexpr double TOLERANCE = 1e-8;
+
+  TwoJetTestHelper<TwoJetL> helper;
+  std::mt19937 rng{SEED};
+  auto dist_mean = helper.make_test_two_jet();
+  const auto samples =
+      random_samples(NUM_SAMPLES, dist_mean, STDDEV, InOut{rng});
+
+  // ACTION
+  const auto sample_mean_sv = two_jet_mean(samples, TOLERANCE, MAX_ITERATIONS);
+
+  // VERIFICATION
+  // Shouldn't converge because MAX_ITERATIONS is 1
+  EXPECT_FALSE(sample_mean_sv.ok());
+}
+
+TEST(TwoJetMeanTest, TestOpposingTwoJets) {
+  // SETUP
+  constexpr int MAX_ITERATIONS = 50;
+  constexpr double TOLERANCE = 1e-6;
+
+  TwoJetTestHelper<TwoJetL> helper;
+  std::vector<TwoJetL> samples;
+  samples.push_back(helper.make_test_two_jet());
+
+  const SE3 &frame_from_ref{samples.front().frame_from_ref()};
+  const SE3 opposing_frame_from_ref{
+      transforms::SO3::exp(M_PI * Vec3::UnitZ()) * frame_from_ref.rotation(),
+      frame_from_ref.translation()};
+
+  samples.emplace_back(
+      opposing_frame_from_ref,
+      samples.front().d_frame_from_ref(),
+      samples.front().d2_frame_from_ref());
+
+  // ACTION
+  const auto sample_mean_sv = two_jet_mean(samples, TOLERANCE, MAX_ITERATIONS);
+
+  // VERIFICATION
+  EXPECT_TRUE(sample_mean_sv.ok());
+}
+
+}  // namespace resim::curves::learning
diff --git a/resim/curves/optimization/BUILD b/resim/curves/optimization/BUILD
index 7769767a..96c207e8 100644
--- a/resim/curves/optimization/BUILD
+++ b/resim/curves/optimization/BUILD
@@ -10,6 +10,7 @@ cc_library(
     name = "two_jet_tangent_space",
     srcs = ["two_jet_tangent_space.cc"],
     hdrs = ["two_jet_tangent_space.hh"],
+    visibility = ["//resim/curves:__subpackages__"],
     deps = [
         "//resim/curves:two_jet",
         "//resim/math:vector_partition",


Reviews:
user: login: iainjwhiteside
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T19:31:22Z
user: login: mikebauer
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T22:20:05Z
user: login: mikebauer
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T22:21:08Z

Review comments:
diff_hunk: @@ -0,0 +1,36 @@
+// Copyright 2025 ReSim, Inc.
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file or at
+// https://opensource.org/licenses/MIT.
+
+#pragma once
+
+#include <span>
+
+#include "resim/curves/two_jet.hh"
+#include "resim/transforms/se3.hh"
+#include "resim/utils/status_value.hh"
+
+namespace resim::curves::learning {
+
+// Compute the mean of a set of two jets. Since two jets exist on a manifold,
+// this function computes the two jet u such that sum_i(x_i (-) u) = 0 where (-)
+// is defined as the left minus operator as defined in
+// resim/curves/optimization/two_jet_tangent_space.hh. Another way of saying
+// this is that u is the mean of all the samples when they're expressed in u's
+// tangent space. We use iteration to find u  and therefore a tolerance and a
path: resim/curves/learning/two_jet_mean.hh
user: login: iainjwhiteside
body: ```suggestion
// tangent space. We use iteration to find u and therefore a tolerance and a
```
created_at: 2025-03-05T19:23:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,44 @@
+// Copyright 2025 ReSim, Inc.
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file or at
+// https://opensource.org/licenses/MIT.
+
+#include "resim/curves/learning/two_jet_mean.hh"
+
+#include "resim/assert/assert.hh"
+#include "resim/curves/optimization/two_jet_tangent_space.hh"
+#include "resim/curves/two_jet.hh"
+
+namespace resim::curves::learning {
+
+using transforms::SE3;
+using TwoJetL = TwoJetL<SE3>;
+using TangentVector = optimization::TwoJetTangentVector<SE3>;
+
+StatusValue<TwoJetL> two_jet_mean(
+    const std::span<const TwoJetL> &samples,
+    const double tolerance,
+    const int max_iterations) {
+  REASSERT(not samples.empty(), "Sample size must be greater than 1");
path: resim/curves/learning/two_jet_mean.cc
user: login: iainjwhiteside
body: ```suggestion
  REASSERT(not samples.empty(), "Sample size must be greater than 1");
  REASSERT(max_iterations > 0, "Number of iterations must be greater than 0");
```

maybe?
created_at: 2025-03-05T19:30:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,36 @@
+// Copyright 2025 ReSim, Inc.
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file or at
+// https://opensource.org/licenses/MIT.
+
+#pragma once
+
+#include <span>
+
+#include "resim/curves/two_jet.hh"
+#include "resim/transforms/se3.hh"
+#include "resim/utils/status_value.hh"
+
+namespace resim::curves::learning {
+
+// Compute the mean of a set of two jets. Since two jets exist on a manifold,
+// this function computes the two jet u such that sum_i(x_i (-) u) = 0 where (-)
+// is defined as the left minus operator as defined in
+// resim/curves/optimization/two_jet_tangent_space.hh. Another way of saying
+// this is that u is the mean of all the samples when they're expressed in u's
+// tangent space. We use iteration to find u  and therefore a tolerance and a
path: resim/curves/learning/two_jet_mean.hh
user: login: mikebauer
body: Good catch!
created_at: 2025-03-05T22:20:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,44 @@
+// Copyright 2025 ReSim, Inc.
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file or at
+// https://opensource.org/licenses/MIT.
+
+#include "resim/curves/learning/two_jet_mean.hh"
+
+#include "resim/assert/assert.hh"
+#include "resim/curves/optimization/two_jet_tangent_space.hh"
+#include "resim/curves/two_jet.hh"
+
+namespace resim::curves::learning {
+
+using transforms::SE3;
+using TwoJetL = TwoJetL<SE3>;
+using TangentVector = optimization::TwoJetTangentVector<SE3>;
+
+StatusValue<TwoJetL> two_jet_mean(
+    const std::span<const TwoJetL> &samples,
+    const double tolerance,
+    const int max_iterations) {
+  REASSERT(not samples.empty(), "Sample size must be greater than 1");
path: resim/curves/learning/two_jet_mean.cc
user: login: mikebauer
body: Good suggestion. Added!
created_at: 2025-03-05T22:21:08Z
author_association: CONTRIBUTOR

PR comments:
user: login: codecov[bot]
created_at: 2025-03-05T22:38:26Z
body: ## [Codecov](https://app.codecov.io/gh/resim-ai/open-core/pull/240?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=resim-ai) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 100.00%. Comparing base [(`d509b46`)](https://app.codecov.io/gh/resim-ai/open-core/commit/d509b46cf47797a6654d4bda730d7e7144c97932?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=resim-ai) to head [(`4b3b727`)](https://app.codecov.io/gh/resim-ai/open-core/commit/4b3b727e33593c9352c18049660130f02eb12708?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=resim-ai).
> Report is 2 commits behind head on main.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff            @@
##              main      #240   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files          254       257    +3     
  Lines        11301     11398   +97     
  Branches      1230      1238    +8     
=========================================
+ Hits         11301     11398   +97     
```

| [Files with missing lines](https://app.codecov.io/gh/resim-ai/open-core/pull/240?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=resim-ai) | Coverage Δ | |
|---|---|---|
| [resim/curves/learning/two\_jet\_mean.cc](https://app.codecov.io/gh/resim-ai/open-core/pull/240?src=pr&el=tree&filepath=resim%2Fcurves%2Flearning%2Ftwo_jet_mean.cc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=resim-ai#diff-cmVzaW0vY3VydmVzL2xlYXJuaW5nL3R3b19qZXRfbWVhbi5jYw==) | `100.00% <100.00%> (ø)` | |

</details>
<details><summary>🚀 New features to boost your workflow: </summary>

- ❄ [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE