Skip to content

Commit

Permalink
[drive] Add custom weighted twist pose estimator
Browse files Browse the repository at this point in the history
Currently we just weight drive measurements to 1 so this should be
basically the same as what we had before.

Signed-off-by: Jade Turner <[email protected]>
  • Loading branch information
spacey-sooty committed Jan 27, 2025
1 parent b930ead commit f173f30
Show file tree
Hide file tree
Showing 4 changed files with 164 additions and 19 deletions.
5 changes: 1 addition & 4 deletions src/main/java/org/curtinfrc/frc2025/Robot.java
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@
import org.curtinfrc.frc2025.subsystems.drive.Drive;
import org.curtinfrc.frc2025.subsystems.drive.GyroIO;
import org.curtinfrc.frc2025.subsystems.drive.GyroIOPigeon2;
import org.curtinfrc.frc2025.subsystems.drive.GyroIOSim;
import org.curtinfrc.frc2025.subsystems.drive.ModuleIO;
import org.curtinfrc.frc2025.subsystems.drive.ModuleIOSim;
import org.curtinfrc.frc2025.subsystems.drive.ModuleIOTalonFX;
Expand Down Expand Up @@ -170,9 +169,7 @@ public Robot() {
// Sim robot, instantiate physics sim IO implementations
drive =
new Drive(
new GyroIOSim(
() -> drive.getKinematics(),
() -> drive.getModuleStates()) {}, // work around crash
new GyroIO() {},
new ModuleIOSim(TunerConstants.FrontLeft),
new ModuleIOSim(TunerConstants.FrontRight),
new ModuleIOSim(TunerConstants.BackLeft),
Expand Down
20 changes: 10 additions & 10 deletions src/main/java/org/curtinfrc/frc2025/subsystems/drive/Drive.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -26,7 +26,6 @@
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;
Expand Down Expand Up @@ -73,8 +72,8 @@ 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 final PIDController xController = new PIDController(10.0, 0.0, 0.0);
private final PIDController yController = new PIDController(10.0, 0.0, 0.0);
Expand Down Expand Up @@ -166,6 +165,8 @@ public void periodic() {
lastModulePositions[moduleIndex] = modulePositions[moduleIndex];
}

Logger.recordOutput("ModuleDeltas", moduleDeltas);

// Update gyro angle
if (gyroInputs.connected) {
// Use the real gyro angle
Expand All @@ -177,7 +178,8 @@ public void periodic() {
}

// Apply update
poseEstimator.updateWithTime(sampleTimestamps[i], rawGyroRotation, modulePositions);
poseEstimator.updateWithTime(
modulePositions, rawGyroRotation, sampleTimestamps[i], VecBuilder.fill(1, 1, 1));
}

// Update gyro alert
Expand Down Expand Up @@ -551,7 +553,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. */
Expand All @@ -566,9 +568,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);
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,151 @@
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;
import org.littletonrobotics.junction.Logger;

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();
Logger.recordOutput("PoseEstimator/PreviousWheelPositions", prevWheelPositions);
Logger.recordOutput("PoseEstimator/CurrentWheelPositions", wheelPositions);
Logger.recordOutput("PoseEstimator/Samples", samples.size());
Logger.recordOutput("PoseEstimator/Twist", twist);

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());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,11 @@

import static org.curtinfrc.frc2025.subsystems.vision.VisionConstants.*;

import edu.wpi.first.math.Matrix;
import edu.wpi.first.math.VecBuilder;
import edu.wpi.first.math.Vector;
import edu.wpi.first.math.geometry.Pose2d;
import edu.wpi.first.math.geometry.Pose3d;
import edu.wpi.first.math.geometry.Rotation2d;
import edu.wpi.first.math.numbers.N1;
import edu.wpi.first.math.numbers.N3;
import edu.wpi.first.wpilibj.Alert;
import edu.wpi.first.wpilibj.Alert.AlertType;
Expand Down Expand Up @@ -187,8 +186,6 @@ public void periodic() {
@FunctionalInterface
public static interface PoseEstimateConsumer {
public void accept(
Pose2d visionRobotPoseMeters,
double timestampSeconds,
Matrix<N3, N1> visionMeasurementStdDevs);
Pose2d visionRobotPoseMeters, double timestampSeconds, Vector<N3> visionMeasurementStdDevs);
}
}

0 comments on commit f173f30

Please sign in to comment.