[drive] Add custom weighted twist pose estimator

src/main/java/org/curtinfrc/frc2025/Constants.java

@@ -30,7 +30,7 @@
  * (log replay from a file).
  */
 public final class Constants {
-  public static final RobotType robotType = RobotType.SIMBOT;
+  public static final RobotType robotType = RobotType.DEVBOT;
   public static final double ROBOT_X = 550; // mm
   public static final double ROBOT_Y = 570;
 

src/main/java/org/curtinfrc/frc2025/Robot.java

@@ -30,7 +30,6 @@
 import org.curtinfrc.frc2025.subsystems.ejector.EjectorIOSim;
 import org.curtinfrc.frc2025.subsystems.elevator.Elevator;
 import org.curtinfrc.frc2025.subsystems.elevator.ElevatorIO;
-import org.curtinfrc.frc2025.subsystems.elevator.ElevatorIONeoMaxMotionLaserCAN;
 import org.curtinfrc.frc2025.subsystems.elevator.ElevatorIOSim;
 import org.curtinfrc.frc2025.subsystems.intake.Intake;
 import org.curtinfrc.frc2025.subsystems.intake.IntakeIO;
@@ -158,12 +157,12 @@ public Robot() {
           vision =
               new Vision(
                   drive::addVisionMeasurement,
-                  new VisionIOLimelightGamepiece(camera0Name),
-                  new VisionIOLimelight(camera1Name, drive::getRotation),
-                  new VisionIOQuestNav());
-          elevator = new Elevator(new ElevatorIONeoMaxMotionLaserCAN());
-          intake = new Intake(new IntakeIONEO());
-          ejector = new Ejector(new EjectorIONEO());
+                  new VisionIO() {},
+                  new VisionIO() {},
+                  new VisionIO() {});
+          elevator = new Elevator(new ElevatorIO() {});
+          intake = new Intake(new IntakeIO() {});
+          ejector = new Ejector(new EjectorIO() {});
         }
 
         case SIMBOT -> {

src/main/java/org/curtinfrc/frc2025/subsystems/drive/Drive.java

@@ -11,10 +11,10 @@
 import edu.wpi.first.hal.FRCNetComm.tResourceType;
 import edu.wpi.first.hal.HAL;
 import edu.wpi.first.math.MathUtil;
-import edu.wpi.first.math.Matrix;
+import edu.wpi.first.math.VecBuilder;
+import edu.wpi.first.math.Vector;
 import edu.wpi.first.math.controller.PIDController;
 import edu.wpi.first.math.controller.ProfiledPIDController;
-import edu.wpi.first.math.estimator.SwerveDrivePoseEstimator;
 import edu.wpi.first.math.filter.SlewRateLimiter;
 import edu.wpi.first.math.geometry.Pose2d;
 import edu.wpi.first.math.geometry.Pose3d;
@@ -26,10 +26,10 @@
 import edu.wpi.first.math.kinematics.SwerveDriveKinematics;
 import edu.wpi.first.math.kinematics.SwerveModulePosition;
 import edu.wpi.first.math.kinematics.SwerveModuleState;
-import edu.wpi.first.math.numbers.N1;
 import edu.wpi.first.math.numbers.N3;
 import edu.wpi.first.math.trajectory.TrapezoidProfile;
 import edu.wpi.first.math.util.Units;
+import edu.wpi.first.util.CircularBuffer;
 import edu.wpi.first.wpilibj.Alert;
 import edu.wpi.first.wpilibj.Alert.AlertType;
 import edu.wpi.first.wpilibj.DriverStation;
@@ -50,6 +50,7 @@
 import org.curtinfrc.frc2025.Constants;
 import org.curtinfrc.frc2025.Constants.Mode;
 import org.curtinfrc.frc2025.generated.TunerConstants;
+import org.curtinfrc.frc2025.util.BufferUtil;
 import org.curtinfrc.frc2025.util.RepulsorFieldPlanner;
 import org.littletonrobotics.junction.AutoLogOutput;
 import org.littletonrobotics.junction.Logger;
@@ -73,8 +74,11 @@ public class Drive extends SubsystemBase {
         new SwerveModulePosition(),
         new SwerveModulePosition()
       };
-  private SwerveDrivePoseEstimator poseEstimator =
-      new SwerveDrivePoseEstimator(kinematics, rawGyroRotation, lastModulePositions, Pose2d.kZero);
+  private PoseEstimator poseEstimator =
+      new PoseEstimator(kinematics, lastModulePositions, rawGyroRotation);
+  private CircularBuffer<Double> xAcceleration = new CircularBuffer<Double>(6);
+  private CircularBuffer<Double> yAcceleration = new CircularBuffer<Double>(6);
+  private ChassisSpeeds lastChassisSpeeds = new ChassisSpeeds();
 
   private final PIDController xController = new PIDController(10.0, 0.0, 0.0);
   private final PIDController yController = new PIDController(10.0, 0.0, 0.0);
@@ -171,17 +175,44 @@ public void periodic() {
         // Use the real gyro angle
         rawGyroRotation = gyroInputs.odometryYawPositions[i];
       } else {
-        // Use the angle delta from the kinematics and module deltas
         Twist2d twist = kinematics.toTwist2d(moduleDeltas);
+        // Use the angle delta from the kinematics and module deltas
         rawGyroRotation = rawGyroRotation.plus(new Rotation2d(twist.dtheta));
       }
 
+      var chassisSpeedsDiff = getChassisSpeeds().minus(lastChassisSpeeds);
+      var xAccel = chassisSpeedsDiff.vxMetersPerSecond / 0.02;
+      xAcceleration.addLast(xAccel);
+      var yAccel = chassisSpeedsDiff.vyMetersPerSecond / 0.02;
+      yAcceleration.addLast(yAccel);
+      var avgX = BufferUtil.average(xAcceleration);
+      var avgY = BufferUtil.average(yAcceleration);
+      Logger.recordOutput("FilteredAccelerationX", avgX);
+      Logger.recordOutput("AccelerationX", chassisSpeedsDiff.vxMetersPerSecond / 0.02);
+      Logger.recordOutput("AccelerationY", chassisSpeedsDiff.vyMetersPerSecond / 0.02);
+      Logger.recordOutput(
+          "Difference",
+          Math.abs(chassisSpeedsDiff.vxMetersPerSecond / 0.02 - gyroInputs.xAcceleration));
+      var xDiff = Math.abs(avgX - gyroIO.xAcceleration());
+      Logger.recordOutput("FilteredDifference", xDiff);
+      var yDiff = Math.abs(avgY - gyroIO.yAcceleration());
+      if (xDiff < 1) {
+        xDiff = 1;
+      }
+      if (yDiff < 1) {
+        yDiff = 1;
+      }
       // Apply update
-      poseEstimator.updateWithTime(sampleTimestamps[i], rawGyroRotation, modulePositions);
+      poseEstimator.updateWithTime(
+          modulePositions,
+          rawGyroRotation,
+          sampleTimestamps[i],
+          VecBuilder.fill(1 / xDiff, 1 / yDiff, 1));
     }
 
     // Update gyro alert
     gyroDisconnectedAlert.set(!gyroInputs.connected && Constants.getMode() != Mode.SIM);
+    lastChassisSpeeds = getChassisSpeeds();
   }
 
   /**
@@ -551,7 +582,7 @@ public double getFFCharacterizationVelocity() {
   /** Returns the current odometry pose. */
   @AutoLogOutput(key = "Odometry/Robot")
   public Pose2d getPose() {
-    return poseEstimator.getEstimatedPosition();
+    return poseEstimator.getEstimatedPose();
   }
 
   /** Returns the current odometry rotation. */
@@ -566,9 +597,7 @@ public void setPose(Pose2d pose) {
 
   /** Adds a new timestamped vision measurement. */
   public void addVisionMeasurement(
-      Pose2d visionRobotPoseMeters,
-      double timestampSeconds,
-      Matrix<N3, N1> visionMeasurementStdDevs) {
+      Pose2d visionRobotPoseMeters, double timestampSeconds, Vector<N3> visionMeasurementStdDevs) {
     poseEstimator.addVisionMeasurement(
         visionRobotPoseMeters, timestampSeconds, visionMeasurementStdDevs);
   }

src/main/java/org/curtinfrc/frc2025/subsystems/drive/GyroIO.java

@@ -15,6 +15,7 @@
 
 import edu.wpi.first.math.geometry.Rotation2d;
 import org.littletonrobotics.junction.AutoLog;
+import org.littletonrobotics.junction.AutoLogOutput;
 
 public interface GyroIO {
   @AutoLog
@@ -24,7 +25,19 @@ public static class GyroIOInputs {
     public double yawVelocityRadPerSec = 0.0;
     public double[] odometryYawTimestamps = new double[] {};
     public Rotation2d[] odometryYawPositions = new Rotation2d[] {};
+    public double xAcceleration = 0;
+    public double yAcceleration = 0;
   }
 
   public default void updateInputs(GyroIOInputs inputs) {}
+
+  @AutoLogOutput(key = "Gyro/FilteredXAcceleration")
+  public default double xAcceleration() {
+    return 0;
+  }
+
+  @AutoLogOutput(key = "Gyro/FilteredYAcceleration")
+  public default double yAcceleration() {
+    return 0;
+  }
 }

src/main/java/org/curtinfrc/frc2025/subsystems/drive/GyroIOPigeon2.java

@@ -13,6 +13,7 @@
 
 package org.curtinfrc.frc2025.subsystems.drive;
 
+import static edu.wpi.first.units.Units.MetersPerSecondPerSecond;
 import static org.curtinfrc.frc2025.subsystems.drive.DriveConstants.*;
 
 import com.ctre.phoenix6.BaseStatusSignal;
@@ -24,8 +25,11 @@
 import edu.wpi.first.math.util.Units;
 import edu.wpi.first.units.measure.Angle;
 import edu.wpi.first.units.measure.AngularVelocity;
+import edu.wpi.first.units.measure.LinearAcceleration;
+import edu.wpi.first.util.CircularBuffer;
 import java.util.Queue;
 import org.curtinfrc.frc2025.generated.TunerConstants;
+import org.curtinfrc.frc2025.util.BufferUtil;
 
 /** IO implementation for Pigeon 2. */
 public class GyroIOPigeon2 implements GyroIO {
@@ -34,9 +38,13 @@ public class GyroIOPigeon2 implements GyroIO {
           TunerConstants.DrivetrainConstants.Pigeon2Id,
           TunerConstants.DrivetrainConstants.CANBusName);
   private final StatusSignal<Angle> yaw = pigeon.getYaw();
+  private final StatusSignal<LinearAcceleration> xAcceleration = pigeon.getAccelerationX();
+  private final StatusSignal<LinearAcceleration> yAcceleration = pigeon.getAccelerationY();
   private final Queue<Double> yawPositionQueue;
   private final Queue<Double> yawTimestampQueue;
   private final StatusSignal<AngularVelocity> yawVelocity = pigeon.getAngularVelocityZWorld();
+  private final CircularBuffer<Double> xAccelerationBuffer = new CircularBuffer<>(6);
+  private final CircularBuffer<Double> yAccelerationBuffer = new CircularBuffer<>(6);
 
   public GyroIOPigeon2() {
     pigeon.getConfigurator().apply(new Pigeon2Configuration());
@@ -50,7 +58,9 @@ public GyroIOPigeon2() {
 
   @Override
   public void updateInputs(GyroIOInputs inputs) {
-    inputs.connected = BaseStatusSignal.refreshAll(yaw, yawVelocity).equals(StatusCode.OK);
+    inputs.connected =
+        BaseStatusSignal.refreshAll(yaw, yawVelocity, xAcceleration, yAcceleration)
+            .equals(StatusCode.OK);
     inputs.yawPosition = Rotation2d.fromDegrees(yaw.getValueAsDouble());
     inputs.yawVelocityRadPerSec = Units.degreesToRadians(yawVelocity.getValueAsDouble());
 
@@ -60,7 +70,21 @@ public void updateInputs(GyroIOInputs inputs) {
         yawPositionQueue.stream()
             .map((Double value) -> Rotation2d.fromDegrees(value))
             .toArray(Rotation2d[]::new);
+    inputs.xAcceleration = xAcceleration.getValue().in(MetersPerSecondPerSecond);
+    xAccelerationBuffer.addLast(inputs.xAcceleration);
+    inputs.yAcceleration = yAcceleration.getValue().in(MetersPerSecondPerSecond);
+    yAccelerationBuffer.addLast(inputs.yAcceleration);
     yawTimestampQueue.clear();
     yawPositionQueue.clear();
   }
+
+  @Override
+  public double xAcceleration() {
+    return BufferUtil.average(xAccelerationBuffer);
+  }
+
+  @Override
+  public double yAcceleration() {
+    return BufferUtil.average(yAccelerationBuffer);
+  }
 }

src/main/java/org/curtinfrc/frc2025/subsystems/drive/PoseEstimator.java

@@ -0,0 +1,146 @@
+package org.curtinfrc.frc2025.subsystems.drive;
+
+import edu.wpi.first.math.Vector;
+import edu.wpi.first.math.geometry.Pose2d;
+import edu.wpi.first.math.geometry.Rotation2d;
+import edu.wpi.first.math.geometry.Twist2d;
+import edu.wpi.first.math.kinematics.SwerveDriveKinematics;
+import edu.wpi.first.math.kinematics.SwerveModulePosition;
+import edu.wpi.first.math.numbers.N3;
+import edu.wpi.first.wpilibj.RobotController;
+import java.util.LinkedList;
+import java.util.OptionalInt;
+
+public class PoseEstimator {
+  private static final double kMaxSampleAge = 0.3;
+
+  private final class TimestampedTwist2d extends Twist2d {
+    public final double timestamp;
+
+    public TimestampedTwist2d(double dx, double dy, double dtheta, double timestamp) {
+      super(dx, dy, dtheta);
+      this.timestamp = timestamp;
+    }
+  }
+
+  private final LinkedList<TimestampedTwist2d> samples = new LinkedList<>();
+  private Pose2d rootPose = new Pose2d();
+  private final SwerveModulePosition[] prevWheelPositions;
+  private final SwerveDriveKinematics kinematics;
+  private Rotation2d gyroOffset;
+  private Rotation2d prevAngle;
+
+  public PoseEstimator(
+      SwerveDriveKinematics kinematics,
+      SwerveModulePosition[] wheelPositions,
+      Rotation2d gyroAngle) {
+    this.kinematics = kinematics;
+    prevWheelPositions = kinematics.copy(wheelPositions);
+    gyroOffset = rootPose.getRotation().minus(gyroAngle);
+    prevAngle = rootPose.getRotation();
+  }
+
+  public void resetPosition(
+      Rotation2d gyroAngle, SwerveModulePosition[] wheelPositions, Pose2d pose) {
+    rootPose = pose;
+    prevAngle = pose.getRotation();
+    gyroOffset = pose.getRotation().minus(gyroAngle);
+    kinematics.copyInto(wheelPositions, prevWheelPositions);
+    samples.clear();
+  }
+
+  private OptionalInt indexForTimestamp(double timestamp) {
+    if (samples.isEmpty()) {
+      return OptionalInt.empty();
+    }
+
+    if (samples.getFirst().timestamp > timestamp) {
+      return OptionalInt.empty();
+    }
+    if (samples.getLast().timestamp < timestamp) {
+      return OptionalInt.of(samples.size() - 1);
+    }
+    int low = 0;
+    int high = samples.size() - 1;
+    while (low <= high) {
+      int mid = (low + high) / 2;
+      double midTime = samples.get(mid).timestamp;
+      if (midTime < timestamp) {
+        low = mid + 1;
+      } else if (midTime > timestamp) {
+        high = mid - 1;
+      } else {
+        return OptionalInt.of(mid);
+      }
+    }
+    return OptionalInt.of(low - 1);
+  }
+
+  private Pose2d poseAtIndex(int index) {
+    // switching this over to primitive math would be a good idea
+    Pose2d pose = rootPose;
+    for (int i = 0; i < index; i++) {
+      pose = pose.exp(samples.get(i));
+    }
+    return pose;
+  }
+
+  private void pruneToRoot() {
+    while (!samples.isEmpty()
+        && samples.getFirst().timestamp < RobotController.getTime() - kMaxSampleAge) {
+      rootPose = rootPose.exp(samples.removeFirst());
+    }
+  }
+
+  /**
+   * Adds a sample to the estimator
+   *
+   * @param pose the pose of the robot at the time of the sample
+   * @param timestamp the timestamp of the sample
+   * @param weight the weight of the sample (0.0 to 1.0)
+   */
+  public void addVisionMeasurement(Pose2d pose, double timestamp, Vector<N3> weight) {
+    var opt = indexForTimestamp(timestamp);
+    if (opt.isEmpty()) {
+      // timestamp is before the first sample
+      return;
+    }
+    int index = opt.getAsInt();
+
+    Pose2d lastPose = poseAtIndex(index);
+    Twist2d twist = lastPose.log(pose);
+    samples.add(
+        index,
+        new TimestampedTwist2d(
+            twist.dx * weight.get(0),
+            twist.dy * weight.get(1),
+            twist.dtheta * weight.get(2),
+            timestamp));
+    pruneToRoot();
+  }
+
+  public void updateWithTime(
+      SwerveModulePosition[] wheelPositions,
+      Rotation2d gyroAngle,
+      double timestamp,
+      Vector<N3> weight) {
+    var angle = gyroAngle.plus(gyroOffset);
+    var twist = kinematics.toTwist2d(prevWheelPositions, wheelPositions);
+    twist.dtheta = angle.minus(prevAngle).getRadians();
+
+    samples.add(
+        new TimestampedTwist2d(
+            twist.dx * weight.get(0),
+            twist.dy * weight.get(1),
+            twist.dtheta * weight.get(2),
+            timestamp));
+
+    kinematics.copyInto(wheelPositions, prevWheelPositions);
+    prevAngle = angle;
+    pruneToRoot();
+  }
+
+  public Pose2d getEstimatedPose() {
+    return poseAtIndex(samples.size());
+  }
+}