AprilTag pose estimation

src/main/java/org/team1540/robot2024/Constants.java

@@ -1,6 +1,8 @@
 package org.team1540.robot2024;
 
+import edu.wpi.first.math.geometry.Pose3d;
 import edu.wpi.first.math.geometry.Rotation2d;
+import edu.wpi.first.math.geometry.Rotation3d;
 import edu.wpi.first.math.util.Units;
 
 /**
@@ -62,6 +64,29 @@ public static class Drivetrain {
         public static final double MAX_ANGULAR_SPEED = MAX_LINEAR_SPEED / DRIVE_BASE_RADIUS;
     }
 
+    public static class Vision {
+        public static final String FRONT_CAMERA_NAME = "limelight-front";
+        public static final String REAR_CAMERA_NAME = "limelight-rear";
+
+        // TODO: measure these offsets
+        public static final Pose3d FRONT_CAMERA_POSE = new Pose3d(0, 0, 0.5, new Rotation3d());
+        public static final Pose3d REAR_CAMERA_POSE = new Pose3d(0, 0, 0.5, new Rotation3d(0, 0, Math.PI));
+
+        // TODO: find these values
+        public static final double MAX_VISION_DELAY_SECS = 0.08;
+        public static final double MAX_ACCEPTED_ROT_SPEED_RAD_PER_SEC = 1.0;
+        public static final double MAX_ACCEPTED_LINEAR_SPEED_MPS = 4.0;
+        public static final double MIN_ACCEPTED_NUM_TAGS = 1;
+        public static final double MAX_ACCEPTED_AVG_TAG_DIST_METERS = 8.0;
+        public static final double MAX_ACCEPTED_ELEVATOR_SPEED_MPS = 0.05;
+
+        public static final int SIM_RES_WIDTH = 1280;
+        public static final int SIM_RES_HEIGHT = 960;
+        public static final Rotation2d SIM_DIAGONAL_FOV = Rotation2d.fromDegrees(100);
+        public static final double SIM_FPS = 14.5;
+        public static final double SIM_AVG_LATENCY_MS = 67.0;
+    }
+
     public static class Elevator {
         public static final double ELEVATOR_MAX_HEIGHT = Units.inchesToMeters(21.0);
     }

src/main/java/org/team1540/robot2024/RobotContainer.java

@@ -14,9 +14,14 @@
 import org.team1540.robot2024.commands.SwerveDriveCommand;
 import org.team1540.robot2024.subsystems.drive.*;
 import org.team1540.robot2024.subsystems.shooter.*;
+import org.team1540.robot2024.subsystems.vision.AprilTagVision;
+import org.team1540.robot2024.subsystems.vision.AprilTagVisionIO;
+import org.team1540.robot2024.subsystems.vision.AprilTagVisionIOLimelight;
+import org.team1540.robot2024.subsystems.vision.AprilTagVisionIOSim;
 import org.team1540.robot2024.util.PhoenixTimeSyncSignalRefresher;
 import org.team1540.robot2024.util.swerve.SwerveFactory;
 import org.littletonrobotics.junction.networktables.LoggedDashboardChooser;
+import org.team1540.robot2024.util.vision.VisionPoseAcceptor;
 
 import static org.team1540.robot2024.Constants.SwerveConfig;
 
@@ -30,6 +35,7 @@ public class RobotContainer {
     // Subsystems
     public final Drivetrain drivetrain;
     public final Shooter shooter;
+    public final AprilTagVision aprilTagVision;
 
     // Controller
     public final CommandXboxController driver = new CommandXboxController(0);
@@ -59,6 +65,12 @@ public RobotContainer() {
                                 new ModuleIOTalonFX(SwerveFactory.getModuleMotors(SwerveConfig.BACK_LEFT, SwerveFactory.SwerveCorner.BACK_LEFT), odometrySignalRefresher),
                                 new ModuleIOTalonFX(SwerveFactory.getModuleMotors(SwerveConfig.BACK_RIGHT, SwerveFactory.SwerveCorner.BACK_RIGHT), odometrySignalRefresher));
                 shooter = new Shooter(new ShooterPivotIOTalonFX(), new FlywheelsIOTalonFX());
+                aprilTagVision = new AprilTagVision(
+                        new AprilTagVisionIOLimelight(Constants.Vision.FRONT_CAMERA_NAME, Constants.Vision.FRONT_CAMERA_POSE),
+                        new AprilTagVisionIOLimelight(Constants.Vision.REAR_CAMERA_NAME, Constants.Vision.REAR_CAMERA_POSE),
+                        drivetrain::addVisionMeasurement,
+                        () -> 0.0, // TODO: ACTUALLY GET ELEVATOR HEIGHT HERE
+                        new VisionPoseAcceptor(drivetrain::getChassisSpeeds, () -> 0.0)); // TODO: ACTUALLY GET ELEVATOR VELOCITY HERE
                 break;
 
             case SIM:
@@ -71,6 +83,13 @@ public RobotContainer() {
                                 new ModuleIOSim(),
                                 new ModuleIOSim());
                 shooter = new Shooter(new ShooterPivotIOSim(), new FlywheelsIOSim());
+                aprilTagVision =
+                        new AprilTagVision(
+                                new AprilTagVisionIOSim(Constants.Vision.FRONT_CAMERA_NAME, Constants.Vision.FRONT_CAMERA_POSE, drivetrain::getPose),
+                                new AprilTagVisionIOSim(Constants.Vision.REAR_CAMERA_NAME, Constants.Vision.REAR_CAMERA_POSE, drivetrain::getPose),
+                                ignored -> {},
+                                () -> 0.0, // TODO: ACTUALLY GET ELEVATOR HEIGHT HERE
+                                new VisionPoseAcceptor(drivetrain::getChassisSpeeds, () -> 0.0)); // TODO: ACTUALLY GET ELEVATOR VELOCITY HERE
                 break;
 
             default:
@@ -83,6 +102,13 @@ public RobotContainer() {
                                 new ModuleIO() {},
                                 new ModuleIO() {});
                 shooter = new Shooter(new ShooterPivotIO() {}, new FlywheelsIO() {});
+                aprilTagVision =
+                        new AprilTagVision(
+                                new AprilTagVisionIO() {},
+                                new AprilTagVisionIO() {},
+                                (ignored) -> {},
+                                () -> 0.0,
+                                new VisionPoseAcceptor(drivetrain::getChassisSpeeds, () -> 0.0));
                 break;
         }
 

src/main/java/org/team1540/robot2024/subsystems/drive/Drivetrain.java

@@ -5,10 +5,11 @@
 import com.pathplanner.lib.util.HolonomicPathFollowerConfig;
 import com.pathplanner.lib.util.PathPlannerLogging;
 import com.pathplanner.lib.util.ReplanningConfig;
+import edu.wpi.first.math.VecBuilder;
+import edu.wpi.first.math.estimator.SwerveDrivePoseEstimator;
 import edu.wpi.first.math.geometry.Pose2d;
 import edu.wpi.first.math.geometry.Rotation2d;
 import edu.wpi.first.math.geometry.Translation2d;
-import edu.wpi.first.math.geometry.Twist2d;
 import edu.wpi.first.math.kinematics.ChassisSpeeds;
 import edu.wpi.first.math.kinematics.SwerveDriveKinematics;
 import edu.wpi.first.math.kinematics.SwerveModulePosition;
@@ -18,6 +19,7 @@
 import org.team1540.robot2024.util.LocalADStarAK;
 import org.littletonrobotics.junction.AutoLogOutput;
 import org.littletonrobotics.junction.Logger;
+import org.team1540.robot2024.util.vision.TimestampedVisionPose;
 
 import static org.team1540.robot2024.Constants.Drivetrain.*;
 
@@ -27,10 +29,11 @@ public class Drivetrain extends SubsystemBase {
     private final Module[] modules = new Module[4]; // FL, FR, BL, BR
 
     private final SwerveDriveKinematics kinematics = new SwerveDriveKinematics(getModuleTranslations());
-    private Pose2d pose = new Pose2d();
-    private Rotation2d lastGyroRotation = new Rotation2d();
+    private Rotation2d rawGyroRotation = new Rotation2d();
     private boolean forceModuleAngleChange = false;
 
+    private final SwerveDrivePoseEstimator poseEstimator;
+
     public Drivetrain(
             GyroIO gyroIO,
             ModuleIO flModuleIO,
@@ -43,6 +46,15 @@ public Drivetrain(
         modules[2] = new Module(blModuleIO, 2);
         modules[3] = new Module(brModuleIO, 3);
 
+        poseEstimator = new SwerveDrivePoseEstimator(
+                kinematics,
+                rawGyroRotation,
+                getModulePositions(),
+                new Pose2d(),
+                // TODO: tune std devs (scale vision by distance?)
+                VecBuilder.fill(0.1, 0.1, 0.1),
+                VecBuilder.fill(0.5, 0.5, 5.0)); // Trust the gyro more than the AprilTags
+
         // Configure AutoBuilder for PathPlanner
         AutoBuilder.configureHolonomic(
                 this::getPose,
@@ -60,44 +72,35 @@ public Drivetrain(
                 (targetPose) -> Logger.recordOutput("Odometry/TrajectorySetpoint", targetPose));
     }
 
+    @Override
     public void periodic() {
         gyroIO.updateInputs(gyroInputs);
         Logger.processInputs("Drivetrain/Gyro", gyroInputs);
         for (Module module : modules) {
             module.periodic();
         }
 
-        // Stop moving when disabled
-        if (DriverStation.isDisabled()) {
-            for (Module module : modules) {
-                module.stop();
-            }
-        }
-        // Log empty setpoint states when disabled
         if (DriverStation.isDisabled()) {
+            // Stop moving when disabled
+            for (Module module : modules) module.stop();
+
+            // Log empty setpoint states when disabled
             Logger.recordOutput("SwerveStates/Setpoints", new SwerveModuleState[]{});
             Logger.recordOutput("SwerveStates/SetpointsOptimized", new SwerveModuleState[]{});
         }
 
-        // Update odometry
+        // Calculate module deltas
         SwerveModulePosition[] wheelDeltas = new SwerveModulePosition[4];
         for (int i = 0; i < 4; i++) {
             wheelDeltas[i] = modules[i].getPositionDelta();
         }
-        // The twist represents the motion of the robot since the last
-        // loop cycle in x, y, and theta based only on the modules,
-        // without the gyro. The gyro is always disconnected in simulation.
-        Twist2d twist = kinematics.toTwist2d(wheelDeltas);
-        if (gyroInputs.connected) {
-            // If the gyro is connected, replace the theta component of the twist
-            // with the change in angle since the last loop cycle.
-            twist =
-                    new Twist2d(
-                            twist.dx, twist.dy, gyroInputs.yawPosition.minus(lastGyroRotation).getRadians());
-            lastGyroRotation = gyroInputs.yawPosition;
-        }
-        // Apply the twist (change since last loop cycle) to the current pose
-        pose = pose.exp(twist);
+        // Use gyro rotation if gyro is connected, otherwise use module deltas to calculate rotation delta
+        rawGyroRotation =
+                gyroInputs.connected ?
+                        gyroInputs.yawPosition
+                        : rawGyroRotation.plus(Rotation2d.fromRadians(kinematics.toTwist2d(wheelDeltas).dtheta));
+        // Update odometry
+        poseEstimator.update(rawGyroRotation, getModulePositions());
     }
 
     /**
@@ -164,6 +167,10 @@ public double getCharacterizationVelocity() {
         return driveVelocityAverage / 4.0;
     }
 
+    public ChassisSpeeds getChassisSpeeds() {
+        return kinematics.toChassisSpeeds(getModuleStates());
+    }
+
     /**
      * Returns the module states (turn angles and drive velocities) for all the modules.
      */
@@ -181,21 +188,34 @@ private SwerveModuleState[] getModuleStates() {
      */
     @AutoLogOutput(key = "Odometry/Robot")
     public Pose2d getPose() {
-        return pose;
+        return poseEstimator.getEstimatedPosition();
     }
 
     /**
      * Returns the current odometry rotation.
      */
     public Rotation2d getRotation() {
-        return pose.getRotation();
+        return getPose().getRotation();
     }
 
     /**
      * Resets the current odometry pose.
      */
     public void setPose(Pose2d pose) {
-        this.pose = pose;
+        poseEstimator.resetPosition(rawGyroRotation, getModulePositions(), pose);
+    }
+
+    public void addVisionMeasurement(TimestampedVisionPose visionPose) {
+        poseEstimator.addVisionMeasurement(visionPose.poseMeters(), visionPose.timestampSecs());
+    }
+
+    public SwerveModulePosition[] getModulePositions() {
+        return new SwerveModulePosition[]{
+                modules[0].getPosition(),
+                modules[1].getPosition(),
+                modules[2].getPosition(),
+                modules[3].getPosition(),
+        };
     }
 
     /**

src/main/java/org/team1540/robot2024/subsystems/vision/AprilTagVision.java

@@ -0,0 +1,119 @@
+package org.team1540.robot2024.subsystems.vision;
+
+import edu.wpi.first.math.geometry.Pose2d;
+import edu.wpi.first.math.geometry.Pose3d;
+import edu.wpi.first.wpilibj2.command.SubsystemBase;
+import org.littletonrobotics.junction.Logger;
+import org.team1540.robot2024.util.vision.TimestampedVisionPose;
+import org.team1540.robot2024.util.vision.VisionPoseAcceptor;
+
+import java.util.Arrays;
+import java.util.Optional;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
+
+import static org.team1540.robot2024.Constants.Vision.*;
+
+public class AprilTagVision extends SubsystemBase {
+    private final AprilTagVisionIO frontCameraIO;
+    private final AprilTagVisionIOInputsAutoLogged frontCameraInputs = new AprilTagVisionIOInputsAutoLogged();
+
+    private final AprilTagVisionIO rearCameraIO;
+    private final AprilTagVisionIOInputsAutoLogged rearCameraInputs = new AprilTagVisionIOInputsAutoLogged();
+
+    private final Consumer<TimestampedVisionPose> visionPoseConsumer;
+    private final Supplier<Double> elevatorHeightSupplierMeters;
+    private final VisionPoseAcceptor poseAcceptor;
+
+    private TimestampedVisionPose frontPose =
+            new TimestampedVisionPose(-1, new Pose2d(), new int[0], new Pose2d[0], true, false);
+    private TimestampedVisionPose rearPose =
+            new TimestampedVisionPose(-1, new Pose2d(), new int[0], new Pose2d[0], false, true);
+
+    public AprilTagVision(
+            AprilTagVisionIO frontCameraIO,
+            AprilTagVisionIO rearCameraIO,
+            Consumer<TimestampedVisionPose> visionPoseConsumer,
+            Supplier<Double> elevatorHeightSupplierMeters,
+            VisionPoseAcceptor poseAcceptor) {
+        this.frontCameraIO = frontCameraIO;
+        this.rearCameraIO = rearCameraIO;
+        this.visionPoseConsumer = visionPoseConsumer;
+        this.elevatorHeightSupplierMeters = elevatorHeightSupplierMeters;
+        this.poseAcceptor = poseAcceptor;
+    }
+
+    @Override
+    public void periodic() {
+        // TODO: need to change if one camera is stationary
+        frontCameraIO.setPoseOffset(
+                new Pose3d(
+                        FRONT_CAMERA_POSE.getX(),
+                        FRONT_CAMERA_POSE.getY(),
+                        FRONT_CAMERA_POSE.getZ() + elevatorHeightSupplierMeters.get(),
+                        FRONT_CAMERA_POSE.getRotation())
+        );
+        rearCameraIO.setPoseOffset(
+                new Pose3d(
+                        REAR_CAMERA_POSE.getX(),
+                        REAR_CAMERA_POSE.getY(),
+                        REAR_CAMERA_POSE.getZ() + elevatorHeightSupplierMeters.get(),
+                        REAR_CAMERA_POSE.getRotation())
+        );
+
+        frontCameraIO.updateInputs(frontCameraInputs);
+        rearCameraIO.updateInputs(rearCameraInputs);
+        Logger.processInputs("Vision/FrontCamera", frontCameraInputs);
+        Logger.processInputs("Vision/RearCamera", rearCameraInputs);
+
+        if (frontCameraInputs.lastMeasurementTimestampSecs > frontPose.timestampSecs()) {
+            frontPose = new TimestampedVisionPose(
+                    frontCameraInputs.lastMeasurementTimestampSecs,
+                    frontCameraInputs.estimatedPoseMeters,
+                    frontCameraInputs.seenTagIDs,
+                    frontCameraInputs.tagPosesMeters,
+                    true,
+                    false);
+        }
+        if (rearCameraInputs.lastMeasurementTimestampSecs > rearPose.timestampSecs()) {
+            rearPose = new TimestampedVisionPose(
+                    rearCameraInputs.lastMeasurementTimestampSecs,
+                    rearCameraInputs.estimatedPoseMeters,
+                    rearCameraInputs.seenTagIDs,
+                    rearCameraInputs.tagPosesMeters,
+                    false,
+                    true);
+        }
+
+        Optional<TimestampedVisionPose> latestPose = getEstimatedPose();
+        latestPose.ifPresent(visionPose -> Logger.recordOutput("Vision/EstimatedPose", visionPose.poseMeters()));
+        latestPose.ifPresent(visionPoseConsumer);
+    }
+
+    /**
+     * Gets the estimated pose by fusing individual computed poses from each camera.
+     * Returns null if no tags seen, in simulation, or if the elevator is moving
+     * too fast
+     */
+    public Optional<TimestampedVisionPose> getEstimatedPose() {
+        boolean useFrontPose = poseAcceptor.shouldAcceptVision(frontPose);
+        boolean useRearPose = poseAcceptor.shouldAcceptVision(rearPose);
+
+        if (useFrontPose && useRearPose) {
+            int[] allTagIDs = Arrays.copyOf(frontPose.seenTagIDs(), frontPose.getNumTagsSeen() + rearPose.getNumTagsSeen());
+            System.arraycopy(rearPose.seenTagIDs(), 0, allTagIDs, frontPose.getNumTagsSeen(), rearPose.getNumTagsSeen());
+            Pose2d[] allTagPoses = Arrays.copyOf(frontPose.tagPosesMeters(), frontPose.getNumTagsSeen() + rearPose.getNumTagsSeen());
+            System.arraycopy(rearPose.tagPosesMeters(), 0, allTagPoses, frontPose.getNumTagsSeen(), rearPose.getNumTagsSeen());
+
+            return Optional.of(new TimestampedVisionPose(
+                    (frontPose.timestampSecs() + rearPose.timestampSecs()) / 2,
+                    frontPose.poseMeters().interpolate(rearPose.poseMeters(), 0.5),
+                    allTagIDs,
+                    allTagPoses,
+                    true,
+                    true));
+        } else if (useFrontPose) return Optional.of(frontPose);
+        else if (useRearPose) return Optional.of(rearPose);
+        else return Optional.empty();
+    }
+}