[wpimath] Add PoseEstimator.sampleAt() (wpilibsuite#6426)

chauser · May 30, 2024 · ff45da3 · ff45da3
1 parent 0981476
commit ff45da3
Show file tree

Hide file tree

Showing 10 changed files with 404 additions and 1 deletion.
diff --git a/wpimath/src/main/java/edu/wpi/first/math/estimator/PoseEstimator.java b/wpimath/src/main/java/edu/wpi/first/math/estimator/PoseEstimator.java
@@ -21,6 +21,7 @@
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Objects;
+import java.util.Optional;
 
 /**
  * This class wraps {@link Odometry} to fuse latency-compensated vision measurements with encoder
@@ -125,6 +126,16 @@ public Pose2d getEstimatedPosition() {
     return m_odometry.getPoseMeters();
   }
 
+  /**
+   * Return the pose at a given timestamp, if the buffer is not empty.
+   *
+   * @param timestampSeconds The pose's timestamp in seconds.
+   * @return The pose at the given timestamp (or Optional.empty() if the buffer is empty).
+   */
+  public Optional<Pose2d> sampleAt(double timestampSeconds) {
+    return m_poseBuffer.getSample(timestampSeconds).map(record -> record.poseMeters);
+  }
+
   /**
    * Adds a vision measurement to the Kalman Filter. This will correct the odometry pose estimate
    * while still accounting for measurement noise.

diff --git a/wpimath/src/main/native/include/frc/estimator/PoseEstimator.h b/wpimath/src/main/native/include/frc/estimator/PoseEstimator.h
@@ -4,6 +4,8 @@
 
 #pragma once
 
+#include <optional>
+
 #include <Eigen/Core>
 #include <wpi/SymbolExports.h>
 #include <wpi/array.h>
@@ -89,6 +91,15 @@ class WPILIB_DLLEXPORT PoseEstimator {
    */
   Pose2d GetEstimatedPosition() const;
 
+  /**
+   * Return the pose at a given timestamp, if the buffer is not empty.
+   *
+   * @param timestamp The pose's timestamp.
+   * @return The pose at the given timestamp (or std::nullopt if the buffer is
+   * empty).
+   */
+  std::optional<Pose2d> SampleAt(units::second_t timestamp) const;
+
   /**
    * Adds a vision measurement to the Kalman Filter. This will correct
    * the odometry pose estimate while still accounting for measurement noise.

diff --git a/wpimath/src/main/native/include/frc/estimator/PoseEstimator.inc b/wpimath/src/main/native/include/frc/estimator/PoseEstimator.inc
@@ -57,6 +57,19 @@ Pose2d PoseEstimator<WheelSpeeds, WheelPositions>::GetEstimatedPosition()
   return m_odometry.GetPose();
 }
 
+template <typename WheelSpeeds, WheelPositions WheelPositions>
+std::optional<Pose2d> PoseEstimator<WheelSpeeds, WheelPositions>::SampleAt(
+    units::second_t timestamp) const {
+  // TODO Replace with std::optional::transform() in C++23
+  std::optional<PoseEstimator<WheelSpeeds, WheelPositions>::InterpolationRecord>
+      sample = m_poseBuffer.Sample(timestamp);
+  if (sample) {
+    return sample->pose;
+  } else {
+    return std::nullopt;
+  }
+}
+
 template <typename WheelSpeeds, WheelPositions WheelPositions>
 void PoseEstimator<WheelSpeeds, WheelPositions>::AddVisionMeasurement(
     const Pose2d& visionRobotPose, units::second_t timestamp) {

diff --git a/wpimath/src/main/native/include/frc/interpolation/TimeInterpolatableBuffer.h b/wpimath/src/main/native/include/frc/interpolation/TimeInterpolatableBuffer.h
@@ -100,7 +100,7 @@ class TimeInterpolatableBuffer {
    *
    * @param time The time at which to sample the buffer.
    */
-  std::optional<T> Sample(units::second_t time) {
+  std::optional<T> Sample(units::second_t time) const {
     if (m_pastSnapshots.empty()) {
       return {};
     }

diff --git a/wpimath/src/test/java/edu/wpi/first/math/estimator/DifferentialDrivePoseEstimatorTest.java b/wpimath/src/test/java/edu/wpi/first/math/estimator/DifferentialDrivePoseEstimatorTest.java
@@ -18,6 +18,7 @@
 import edu.wpi.first.math.trajectory.TrajectoryConfig;
 import edu.wpi.first.math.trajectory.TrajectoryGenerator;
 import java.util.List;
+import java.util.Optional;
 import java.util.Random;
 import java.util.TreeMap;
 import java.util.function.Function;
@@ -304,4 +305,54 @@ void testDiscardsStaleVisionMeasurements() {
         estimator.getEstimatedPosition().getRotation().getRadians(),
         "Incorrect Final Theta");
   }
+
+  @Test
+  void testSampleAt() {
+    var kinematics = new DifferentialDriveKinematics(1);
+    var estimator =
+        new DifferentialDrivePoseEstimator(
+            kinematics,
+            new Rotation2d(),
+            0,
+            0,
+            new Pose2d(),
+            VecBuilder.fill(1, 1, 1),
+            VecBuilder.fill(1, 1, 1));
+
+    // Returns empty when null
+    assertEquals(Optional.empty(), estimator.sampleAt(1));
+
+    // Add odometry measurements, but don't fill up the buffer
+    // Add a tiny tolerance for the upper bound because of floating point rounding error
+    for (double time = 1; time <= 2 + 1e-9; time += 0.02) {
+      estimator.updateWithTime(time, new Rotation2d(), time, time);
+    }
+
+    // Sample at an added time
+    assertEquals(Optional.of(new Pose2d(1.02, 0, new Rotation2d())), estimator.sampleAt(1.02));
+    // Sample between updates (test interpolation)
+    assertEquals(Optional.of(new Pose2d(1.01, 0, new Rotation2d())), estimator.sampleAt(1.01));
+    // Sampling before the oldest value returns the oldest value
+    assertEquals(Optional.of(new Pose2d(1, 0, new Rotation2d())), estimator.sampleAt(0.5));
+    // Sampling after the newest value returns the newest value
+    assertEquals(Optional.of(new Pose2d(2, 0, new Rotation2d())), estimator.sampleAt(2.5));
+
+    // Add a vision measurement after the odometry measurements (while keeping all of the old
+    // odometry measurements)
+    estimator.addVisionMeasurement(new Pose2d(2, 0, new Rotation2d(1)), 2.2);
+
+    // Make sure nothing changed (except the newest value)
+    assertEquals(Optional.of(new Pose2d(1.02, 0, new Rotation2d())), estimator.sampleAt(1.02));
+    assertEquals(Optional.of(new Pose2d(1.01, 0, new Rotation2d())), estimator.sampleAt(1.01));
+    assertEquals(Optional.of(new Pose2d(1, 0, new Rotation2d())), estimator.sampleAt(0.5));
+
+    // Add a vision measurement before the odometry measurements that's still in the buffer
+    estimator.addVisionMeasurement(new Pose2d(1, 0.2, new Rotation2d()), 0.9);
+
+    // Everything should be the same except Y is 0.1 (halfway between 0 and 0.2)
+    assertEquals(Optional.of(new Pose2d(1.02, 0.1, new Rotation2d())), estimator.sampleAt(1.02));
+    assertEquals(Optional.of(new Pose2d(1.01, 0.1, new Rotation2d())), estimator.sampleAt(1.01));
+    assertEquals(Optional.of(new Pose2d(1, 0.1, new Rotation2d())), estimator.sampleAt(0.5));
+    assertEquals(Optional.of(new Pose2d(2, 0.1, new Rotation2d())), estimator.sampleAt(2.5));
+  }
 }
diff --git a/wpimath/src/test/java/edu/wpi/first/math/estimator/MecanumDrivePoseEstimatorTest.java b/wpimath/src/test/java/edu/wpi/first/math/estimator/MecanumDrivePoseEstimatorTest.java
@@ -19,6 +19,7 @@
 import edu.wpi.first.math.trajectory.TrajectoryConfig;
 import edu.wpi.first.math.trajectory.TrajectoryGenerator;
 import java.util.List;
+import java.util.Optional;
 import java.util.Random;
 import java.util.TreeMap;
 import java.util.function.Function;
@@ -321,4 +322,59 @@ void testDiscardsOldVisionMeasurements() {
         estimator.getEstimatedPosition().getRotation().getRadians(),
         "Incorrect Final Theta");
   }
+
+  @Test
+  void testSampleAt() {
+    var kinematics =
+        new MecanumDriveKinematics(
+            new Translation2d(1, 1),
+            new Translation2d(-1, 1),
+            new Translation2d(1, -1),
+            new Translation2d(-1, -1));
+    var estimator =
+        new MecanumDrivePoseEstimator(
+            kinematics,
+            new Rotation2d(),
+            new MecanumDriveWheelPositions(),
+            new Pose2d(),
+            VecBuilder.fill(1, 1, 1),
+            VecBuilder.fill(1, 1, 1));
+
+    // Returns empty when null
+    assertEquals(Optional.empty(), estimator.sampleAt(1));
+
+    // Add odometry measurements, but don't fill up the buffer
+    // Add a tiny tolerance for the upper bound because of floating point rounding error
+    for (double time = 1; time <= 2 + 1e-9; time += 0.02) {
+      var wheelPositions = new MecanumDriveWheelPositions(time, time, time, time);
+      estimator.updateWithTime(time, new Rotation2d(), wheelPositions);
+    }
+
+    // Sample at an added time
+    assertEquals(Optional.of(new Pose2d(1.02, 0, new Rotation2d())), estimator.sampleAt(1.02));
+    // Sample between updates (test interpolation)
+    assertEquals(Optional.of(new Pose2d(1.01, 0, new Rotation2d())), estimator.sampleAt(1.01));
+    // Sampling before the oldest value returns the oldest value
+    assertEquals(Optional.of(new Pose2d(1, 0, new Rotation2d())), estimator.sampleAt(0.5));
+    // Sampling after the newest value returns the newest value
+    assertEquals(Optional.of(new Pose2d(2, 0, new Rotation2d())), estimator.sampleAt(2.5));
+
+    // Add a vision measurement after the odometry measurements (while keeping all of the old
+    // odometry measurements)
+    estimator.addVisionMeasurement(new Pose2d(2, 0, new Rotation2d(1)), 2.2);
+
+    // Make sure nothing changed (except the newest value)
+    assertEquals(Optional.of(new Pose2d(1.02, 0, new Rotation2d())), estimator.sampleAt(1.02));
+    assertEquals(Optional.of(new Pose2d(1.01, 0, new Rotation2d())), estimator.sampleAt(1.01));
+    assertEquals(Optional.of(new Pose2d(1, 0, new Rotation2d())), estimator.sampleAt(0.5));
+
+    // Add a vision measurement before the odometry measurements that's still in the buffer
+    estimator.addVisionMeasurement(new Pose2d(1, 0.2, new Rotation2d()), 0.9);
+
+    // Everything should be the same except Y is 0.1 (halfway between 0 and 0.2)
+    assertEquals(Optional.of(new Pose2d(1.02, 0.1, new Rotation2d())), estimator.sampleAt(1.02));
+    assertEquals(Optional.of(new Pose2d(1.01, 0.1, new Rotation2d())), estimator.sampleAt(1.01));
+    assertEquals(Optional.of(new Pose2d(1, 0.1, new Rotation2d())), estimator.sampleAt(0.5));
+    assertEquals(Optional.of(new Pose2d(2, 0.1, new Rotation2d())), estimator.sampleAt(2.5));
+  }
 }
diff --git a/wpimath/src/test/java/edu/wpi/first/math/estimator/SwerveDrivePoseEstimatorTest.java b/wpimath/src/test/java/edu/wpi/first/math/estimator/SwerveDrivePoseEstimatorTest.java
@@ -19,6 +19,7 @@
 import edu.wpi.first.math.trajectory.TrajectoryConfig;
 import edu.wpi.first.math.trajectory.TrajectoryGenerator;
 import java.util.List;
+import java.util.Optional;
 import java.util.Random;
 import java.util.TreeMap;
 import java.util.function.Function;
@@ -355,4 +356,70 @@ void testDiscardsOldVisionMeasurements() {
         estimator.getEstimatedPosition().getRotation().getRadians(),
         "Incorrect Final Theta");
   }
+
+  @Test
+  void testSampleAt() {
+    var kinematics =
+        new SwerveDriveKinematics(
+            new Translation2d(1, 1),
+            new Translation2d(-1, 1),
+            new Translation2d(1, -1),
+            new Translation2d(-1, -1));
+    var estimator =
+        new SwerveDrivePoseEstimator(
+            kinematics,
+            new Rotation2d(),
+            new SwerveModulePosition[] {
+              new SwerveModulePosition(),
+              new SwerveModulePosition(),
+              new SwerveModulePosition(),
+              new SwerveModulePosition()
+            },
+            new Pose2d(),
+            VecBuilder.fill(1, 1, 1),
+            VecBuilder.fill(1, 1, 1));
+
+    // Returns empty when null
+    assertEquals(Optional.empty(), estimator.sampleAt(1));
+
+    // Add odometry measurements, but don't fill up the buffer
+    // Add a tiny tolerance for the upper bound because of floating point rounding error
+    for (double time = 1; time <= 2 + 1e-9; time += 0.02) {
+      var wheelPositions =
+          new SwerveModulePosition[] {
+            new SwerveModulePosition(time, new Rotation2d()),
+            new SwerveModulePosition(time, new Rotation2d()),
+            new SwerveModulePosition(time, new Rotation2d()),
+            new SwerveModulePosition(time, new Rotation2d())
+          };
+      estimator.updateWithTime(time, new Rotation2d(), wheelPositions);
+    }
+
+    // Sample at an added time
+    assertEquals(Optional.of(new Pose2d(1.02, 0, new Rotation2d())), estimator.sampleAt(1.02));
+    // Sample between updates (test interpolation)
+    assertEquals(Optional.of(new Pose2d(1.01, 0, new Rotation2d())), estimator.sampleAt(1.01));
+    // Sampling before the oldest value returns the oldest value
+    assertEquals(Optional.of(new Pose2d(1, 0, new Rotation2d())), estimator.sampleAt(0.5));
+    // Sampling after the newest value returns the newest value
+    assertEquals(Optional.of(new Pose2d(2, 0, new Rotation2d())), estimator.sampleAt(2.5));
+
+    // Add a vision measurement after the odometry measurements (while keeping all of the old
+    // odometry measurements)
+    estimator.addVisionMeasurement(new Pose2d(2, 0, new Rotation2d(1)), 2.2);
+
+    // Make sure nothing changed (except the newest value)
+    assertEquals(Optional.of(new Pose2d(1.02, 0, new Rotation2d())), estimator.sampleAt(1.02));
+    assertEquals(Optional.of(new Pose2d(1.01, 0, new Rotation2d())), estimator.sampleAt(1.01));
+    assertEquals(Optional.of(new Pose2d(1, 0, new Rotation2d())), estimator.sampleAt(0.5));
+
+    // Add a vision measurement before the odometry measurements that's still in the buffer
+    estimator.addVisionMeasurement(new Pose2d(1, 0.2, new Rotation2d()), 0.9);
+
+    // Everything should be the same except Y is 0.1 (halfway between 0 and 0.2)
+    assertEquals(Optional.of(new Pose2d(1.02, 0.1, new Rotation2d())), estimator.sampleAt(1.02));
+    assertEquals(Optional.of(new Pose2d(1.01, 0.1, new Rotation2d())), estimator.sampleAt(1.01));
+    assertEquals(Optional.of(new Pose2d(1, 0.1, new Rotation2d())), estimator.sampleAt(0.5));
+    assertEquals(Optional.of(new Pose2d(2, 0.1, new Rotation2d())), estimator.sampleAt(2.5));
+  }
 }
diff --git a/wpimath/src/test/native/cpp/estimator/DifferentialDrivePoseEstimatorTest.cpp b/wpimath/src/test/native/cpp/estimator/DifferentialDrivePoseEstimatorTest.cpp
@@ -293,3 +293,62 @@ TEST(DifferentialDrivePoseEstimatorTest, TestDiscardStaleVisionMeasurements) {
               estimator.GetEstimatedPosition().Rotation().Radians().value(),
               1e-6);
 }
+
+TEST(DifferentialDrivePoseEstimatorTest, TestSampleAt) {
+  frc::DifferentialDriveKinematics kinematics{1_m};
+  frc::DifferentialDrivePoseEstimator estimator{
+      kinematics,      frc::Rotation2d{}, 0_m, 0_m, frc::Pose2d{},
+      {1.0, 1.0, 1.0}, {1.0, 1.0, 1.0}};
+
+  // Returns empty when null
+  EXPECT_EQ(std::nullopt, estimator.SampleAt(1_s));
+
+  // Add odometry measurements, but don't fill up the buffer
+  // Add a tiny tolerance for the upper bound because of floating point rounding
+  // error
+  for (double time = 1; time <= 2 + 1e-9; time += 0.02) {
+    estimator.UpdateWithTime(units::second_t{time}, frc::Rotation2d{},
+                             units::meter_t{time}, units::meter_t{time});
+  }
+
+  // Sample at an added time
+  EXPECT_EQ(std::optional(frc::Pose2d{1.02_m, 0_m, frc::Rotation2d{}}),
+            estimator.SampleAt(1.02_s));
+  // Sample between updates (test interpolation)
+  EXPECT_EQ(std::optional(frc::Pose2d{1.01_m, 0_m, frc::Rotation2d{}}),
+            estimator.SampleAt(1.01_s));
+  // Sampling before the oldest value returns the oldest value
+  EXPECT_EQ(std::optional(frc::Pose2d{1_m, 0_m, frc::Rotation2d{}}),
+            estimator.SampleAt(0.5_s));
+  // Sampling after the newest value returns the newest value
+  EXPECT_EQ(std::optional(frc::Pose2d{2_m, 0_m, frc::Rotation2d{}}),
+            estimator.SampleAt(2.5_s));
+
+  // Add a vision measurement after the odometry measurements (while keeping all
+  // of the old odometry measurements)
+  estimator.AddVisionMeasurement(frc::Pose2d{2_m, 0_m, frc::Rotation2d{1_rad}},
+                                 2.2_s);
+
+  // Make sure nothing changed (except the newest value)
+  EXPECT_EQ(std::optional(frc::Pose2d{1.02_m, 0_m, frc::Rotation2d{}}),
+            estimator.SampleAt(1.02_s));
+  EXPECT_EQ(std::optional(frc::Pose2d{1.01_m, 0_m, frc::Rotation2d{}}),
+            estimator.SampleAt(1.01_s));
+  EXPECT_EQ(std::optional(frc::Pose2d{1_m, 0_m, frc::Rotation2d{}}),
+            estimator.SampleAt(0.5_s));
+
+  // Add a vision measurement before the odometry measurements that's still in
+  // the buffer
+  estimator.AddVisionMeasurement(frc::Pose2d{1_m, 0.2_m, frc::Rotation2d{}},
+                                 0.9_s);
+
+  // Everything should be the same except Y is 0.1 (halfway between 0 and 0.2)
+  EXPECT_EQ(std::optional(frc::Pose2d{1.02_m, 0.1_m, frc::Rotation2d{}}),
+            estimator.SampleAt(1.02_s));
+  EXPECT_EQ(std::optional(frc::Pose2d{1.01_m, 0.1_m, frc::Rotation2d{}}),
+            estimator.SampleAt(1.01_s));
+  EXPECT_EQ(std::optional(frc::Pose2d{1_m, 0.1_m, frc::Rotation2d{}}),
+            estimator.SampleAt(0.5_s));
+  EXPECT_EQ(std::optional(frc::Pose2d{2_m, 0.1_m, frc::Rotation2d{}}),
+            estimator.SampleAt(2.5_s));
+}