Replace KF in robot localizer with EKF

src/shared/robot_constants.cpp

@@ -53,8 +53,10 @@ RobotConstants createRobotConstants()
 
         // Kalman filter variances for robot localizer
         .kalman_process_noise_variance_rad_per_s_4      = 1.0f,
-        .kalman_vision_noise_variance_rad_2             = 0.01f * 0.01f,
-        .kalman_motor_sensor_noise_variance_rad_per_s_2 = 0.5f};
+        .kalman_vision_noise_variance_rad_2             = 0.0001f,
+        .kalman_motor_sensor_noise_variance_rad_per_s_2 = 0.5f,
+        .kalman_motor_sensor_noise_variance_m_per_s_2 = 0.05f
+	};
 }
 #elif CHECK_VERSION(2021)
 constexpr RobotConstants createRobotConstants()

src/shared/robot_constants.h

@@ -131,6 +131,7 @@ struct RobotConstants
     float kalman_vision_noise_variance_rad_2;
 
     float kalman_motor_sensor_noise_variance_rad_per_s_2;
+    float kalman_motor_sensor_noise_variance_m_per_s_2;
 };
 
 /**

src/software/embedded/BUILD

@@ -39,6 +39,7 @@ cc_library(
         "//proto/primitive:primitive_msg_factory",
         "//software/ai/navigator/trajectory:bang_bang_trajectory_1d_angular",
         "//software/ai/navigator/trajectory:trajectory_path",
+        "//software/embedded:robot_localizer",
         "//software/embedded/motion_control:orientation_controller",
         "//software/embedded/motion_control:position_controller",
         "//software/math:math_functions",
@@ -126,6 +127,8 @@ cc_library(
         "//software/geom:angular_velocity",
         "//software/geom:point",
         "//software/geom:vector",
+        "//software/physics:velocity_conversion_util",
+        "//software/sensor_fusion/filter:extended_kalman_filter",
         "//software/sensor_fusion/filter:kalman_filter",
         "//software/world:robot_state",
         "@eigen",

src/software/embedded/primitive_executor.cpp

@@ -8,13 +8,18 @@
 #include "proto/primitive/primitive_msg_factory.h"
 #include "proto/tbots_software_msgs.pb.h"
 #include "proto/visualization.pb.h"
+#include "software/embedded/robot_localizer.h"
 #include "software/geom/algorithms/distance.h"
 #include "software/logger/logger.h"
 #include "software/physics/velocity_conversion_util.h"
 
 PrimitiveExecutor::PrimitiveExecutor(
-    const robot_constants::RobotConstants& robot_constants, const RobotId robot_id)
-    : robot_state_(), robot_constants_(robot_constants), robot_id_(robot_id)
+    const robot_constants::RobotConstants& robot_constants, const RobotId robot_id,
+    const TeamColour team_colour)
+    : robot_state_(),
+      robot_constants_(robot_constants),
+      robot_id_(robot_id),
+      team_colour_(team_colour)
 {
 }
 

src/software/embedded/primitive_executor.h

@@ -9,6 +9,7 @@
 #include "software/geom/vector.h"
 #include "software/time/duration.h"
 #include "software/world/robot_state.h"
+#include "software/world/team_types.h"
 
 /**
  * "Executes" primitives, turning them into the direct control commands that
@@ -25,9 +26,14 @@ class PrimitiveExecutor
      *
      * @param robot_constants The constants for the robot using this primitive executor
      * @param robot_id The ID of the robot using this primitive executor
+     * @param team_colour The colour of the team this robot belongs to. Only matters for
+     * disambiguating PlotJuggler log keys when multiple teams share one process (e.g.
+     * the simulator, where a yellow and blue robot can have the same ID); real hardware
+     * only ever runs one robot so this can be left at its default.
      */
     explicit PrimitiveExecutor(const robot_constants::RobotConstants& robot_constants,
-                               RobotId robot_id);
+                               RobotId robot_id,
+                               TeamColour team_colour = TeamColour::YELLOW);
 
     /**
      * Starts executing a new primitive.
@@ -117,6 +123,7 @@ class PrimitiveExecutor
 
     robot_constants::RobotConstants robot_constants_;
     RobotId robot_id_;
+    TeamColour team_colour_;
 
     std::optional<TrajectoryPath> trajectory_path_;
     std::optional<BangBangTrajectory1DAngular> angular_trajectory_;

src/software/embedded/robot_localizer.cpp

@@ -1,5 +1,7 @@
 #include "robot_localizer.h"
 
+#include <cmath>
+
 #include "proto/message_translation/tbots_geometry.h"
 #include "shared/constants.h"
 #include "software/physics/velocity_conversion_util.h"
@@ -189,7 +191,12 @@ Point RobotLocalizer::getPosition() const
         filter_.state_estimate(static_cast<Eigen::Index>(StateIndex::Y_POSITION)));
 }
 
-Vector RobotLocalizer::getVelocity() const
+Vector RobotLocalizer::getGlobalVelocity() const
+{
+    return localToGlobalVelocity(getLocalVelocity(), getOrientation());
+}
+
+Vector RobotLocalizer::getLocalVelocity() const
 {
     return Vector(
         filter_.state_estimate(static_cast<Eigen::Index>(StateIndex::X_VELOCITY)),
@@ -211,23 +218,92 @@ AngularVelocity RobotLocalizer::getAngularVelocity() const
 
 RobotState RobotLocalizer::getRobotState() const
 {
-    return RobotState(getPosition(), getVelocity(), getOrientation(),
+    return RobotState(getPosition(), getGlobalVelocity(), getOrientation(),
                       getAngularVelocity());
 }
 
 // TODO: Investigate proces models/variances/etc
 void RobotLocalizer::generatedPredictionMatrices(double delta_time_seconds)
 {
-	// In the current model, we use target velocity as our new velocity of the preiction state, and position is derived from it.
-	// Therefore, process model keeps the positions and we don't predict it using estimated velocities
-    filter_.process_model <<
-        1, 0, 0, 0, 0, 0,
-        0, 1, 0, 0, 0, 0,
-        0, 0, 1, 0, 0, delta_time_seconds,
-        0, 0, 0, 0, 0, 0,
-        0, 0, 0, 0, 0, 0,
-        0, 0, 0, 0, 0, 1;
-    // clang-format on
+    // Velocity is estimated in the robot's local frame (see StateIndex), but position
+    // is in world space, so propagating position requires rotating local velocity by
+    // the current orientation estimate -- a nonlinear operation, hence the process
+    // model function/Jacobian pair instead of a constant matrix.
+    //
+    // Velocity itself isn't propagated from its own estimate: it's replaced outright
+    // by the (rotated) control input every step (see control_model below), so f(x)
+    // leaves it at zero and its row of the Jacobian is zero too.
+    filter_.process_model_function =
+        [delta_time_seconds](Eigen::Vector<double, STATE_SIZE> state)
+    {
+        const double theta =
+            state(static_cast<Eigen::Index>(StateIndex::ORIENTATION));
+        const double local_vx =
+            state(static_cast<Eigen::Index>(StateIndex::X_VELOCITY));
+        const double local_vy =
+            state(static_cast<Eigen::Index>(StateIndex::Y_VELOCITY));
+
+        Eigen::Vector<double, STATE_SIZE> next_state =
+            Eigen::Vector<double, STATE_SIZE>::Zero();
+
+        next_state(static_cast<Eigen::Index>(StateIndex::X_POSITION)) =
+            state(static_cast<Eigen::Index>(StateIndex::X_POSITION)) +
+            delta_time_seconds *
+                (local_vx * std::cos(theta) - local_vy * std::sin(theta));
+        next_state(static_cast<Eigen::Index>(StateIndex::Y_POSITION)) =
+            state(static_cast<Eigen::Index>(StateIndex::Y_POSITION)) +
+            delta_time_seconds *
+                (local_vx * std::sin(theta) + local_vy * std::cos(theta));
+        next_state(static_cast<Eigen::Index>(StateIndex::ORIENTATION)) =
+            theta + delta_time_seconds *
+                        state(static_cast<Eigen::Index>(StateIndex::ANGULAR_VELOCITY));
+        next_state(static_cast<Eigen::Index>(StateIndex::ANGULAR_VELOCITY)) =
+            state(static_cast<Eigen::Index>(StateIndex::ANGULAR_VELOCITY));
+
+        return next_state;
+    };
+
+    filter_.process_model_jacobian_function =
+        [delta_time_seconds](Eigen::Vector<double, STATE_SIZE> state)
+    {
+        const auto x_position_index = static_cast<Eigen::Index>(StateIndex::X_POSITION);
+        const auto y_position_index = static_cast<Eigen::Index>(StateIndex::Y_POSITION);
+        const auto orientation_index =
+            static_cast<Eigen::Index>(StateIndex::ORIENTATION);
+        const auto x_velocity_index = static_cast<Eigen::Index>(StateIndex::X_VELOCITY);
+        const auto y_velocity_index = static_cast<Eigen::Index>(StateIndex::Y_VELOCITY);
+        const auto angular_velocity_index =
+            static_cast<Eigen::Index>(StateIndex::ANGULAR_VELOCITY);
+
+        const double theta       = state(orientation_index);
+        const double local_vx    = state(x_velocity_index);
+        const double local_vy    = state(y_velocity_index);
+        const double cos_theta   = std::cos(theta);
+        const double sin_theta   = std::sin(theta);
+
+        Eigen::Matrix<double, STATE_SIZE, STATE_SIZE> jacobian =
+            Eigen::Matrix<double, STATE_SIZE, STATE_SIZE>::Identity();
+
+        jacobian(x_position_index, orientation_index) =
+            delta_time_seconds * (-local_vx * sin_theta - local_vy * cos_theta);
+        jacobian(x_position_index, x_velocity_index) = delta_time_seconds * cos_theta;
+        jacobian(x_position_index, y_velocity_index) = -delta_time_seconds * sin_theta;
+
+        jacobian(y_position_index, orientation_index) =
+            delta_time_seconds * (local_vx * cos_theta - local_vy * sin_theta);
+        jacobian(y_position_index, x_velocity_index) = delta_time_seconds * sin_theta;
+        jacobian(y_position_index, y_velocity_index) = delta_time_seconds * cos_theta;
+
+        jacobian(orientation_index, angular_velocity_index) = delta_time_seconds;
+
+        // f leaves velocity at zero regardless of the input state (see
+        // process_model_function above), so its row of the Jacobian is zero, not the
+        // identity default.
+        jacobian(x_velocity_index, x_velocity_index) = 0;
+        jacobian(y_velocity_index, y_velocity_index) = 0;
+
+        return jacobian;
+    };
 
     const double delta_time_squared = delta_time_seconds * delta_time_seconds;
     const double delta_time_cubed   = delta_time_squared * delta_time_seconds;
@@ -264,25 +340,33 @@ void RobotLocalizer::generatedPredictionMatrices(double delta_time_seconds)
         0, 0, angular_position_velocity_covariance, 0, 0, angular_velocity_variance;
     // clang-format on
 
-    // Control input is the commanded (target) linear velocity: it replaces the old
-    // velocity state outright (see process_model above) and drives position over this
-    // step's elapsed time.
+    // Control input is the commanded (target) linear velocity in world space: it
+    // replaces the local velocity state outright, rotated into the robot's local
+    // frame by the current orientation estimate (see process_model_function above,
+    // which then rotates that local velocity back into world space to propagate
+    // position). Position is no longer driven directly from control input here --
+    // that happens through the process model function instead.
+    const double theta =
+        filter_.state_estimate(static_cast<Eigen::Index>(StateIndex::ORIENTATION));
+    const double cos_theta = std::cos(theta);
+    const double sin_theta = std::sin(theta);
+
     auto& control_model = filter_.control_model;
     control_model.setZero();
 
-    control_model(static_cast<Eigen::Index>(StateIndex::X_POSITION),
+    control_model(static_cast<Eigen::Index>(StateIndex::X_VELOCITY),
                   static_cast<Eigen::Index>(ControlIndex::X_VELOCITY_TARGET)) =
-        delta_time_seconds;
-
-    control_model(static_cast<Eigen::Index>(StateIndex::Y_POSITION),
-                  static_cast<Eigen::Index>(ControlIndex::Y_VELOCITY_TARGET)) =
-        delta_time_seconds;
-
+        cos_theta;
     control_model(static_cast<Eigen::Index>(StateIndex::X_VELOCITY),
-                  static_cast<Eigen::Index>(ControlIndex::X_VELOCITY_TARGET)) = 1;
+                  static_cast<Eigen::Index>(ControlIndex::Y_VELOCITY_TARGET)) =
+        sin_theta;
 
     control_model(static_cast<Eigen::Index>(StateIndex::Y_VELOCITY),
-                  static_cast<Eigen::Index>(ControlIndex::Y_VELOCITY_TARGET)) = 1;
+                  static_cast<Eigen::Index>(ControlIndex::X_VELOCITY_TARGET)) =
+        -sin_theta;
+    control_model(static_cast<Eigen::Index>(StateIndex::Y_VELOCITY),
+                  static_cast<Eigen::Index>(ControlIndex::Y_VELOCITY_TARGET)) =
+        cos_theta;
 }
 
 void RobotLocalizer::generateMeasurementModel(FilterStepType source)

src/software/embedded/robot_localizer.h

@@ -5,19 +5,23 @@
 #include <optional>
 
 #include "proto/primitive.pb.h"
-#include "proto/robot_status_msg.pb.h"
 #include "software/embedded/services/imu.h"
 #include "software/geom/angle.h"
 #include "software/geom/point.h"
 #include "software/geom/vector.h"
-#include "software/sensor_fusion/filter/kalman_filter.hpp"
+#include "software/sensor_fusion/filter/extended_kalman_filter.hpp"
 #include "software/time/duration.h"
 #include "software/util/make_enum/make_enum.hpp"
 #include "software/world/robot_state.h"
 
+// X_POSITION/Y_POSITION are in world space; X_VELOCITY/Y_VELOCITY are in the robot's
+// local frame (see velocity_conversion_util.h), matching what the motor sensors report
+// directly and avoiding a lossy conversion through the orientation estimate.
 MAKE_ENUM(StateIndex, X_POSITION, Y_POSITION, ORIENTATION, X_VELOCITY, Y_VELOCITY,
           ANGULAR_VELOCITY);
 
+// MOTOR_X_VELOCITY/MOTOR_Y_VELOCITY are in the robot's local frame, matching
+// StateIndex::X_VELOCITY/Y_VELOCITY.
 MAKE_ENUM(MeasurementIndex, VISION_X_POSITION, VISION_Y_POSITION, VISION_ORIENTATION,
           MOTOR_X_VELOCITY, MOTOR_Y_VELOCITY, MOTOR_ANGULAR_VELOCITY,
           IMU_ANGULAR_VELOCITY);
@@ -27,8 +31,12 @@ MAKE_ENUM(ControlIndex, X_VELOCITY_TARGET, Y_VELOCITY_TARGET);
 MAKE_ENUM(FilterStepType, PREDICT, MOTOR_DATA, IMU_DATA, VISION_DATA);
 
 /**
- * Estimates robot orientation, angular velocity, and angular acceleration
- * using a Kalman filter.
+ * Estimates robot position, orientation, velocity, and angular velocity using an
+ * extended Kalman filter.
+ *
+ * The process model is nonlinear because velocity is estimated in the robot's local
+ * frame (see StateIndex) while position is in world space, so propagating position
+ * requires rotating local velocity by the current orientation estimate.
  *
  * The filter keeps a history of recent predict/update operations. When delayed
  * vision data arrives, the localizer rewinds to the matching historical state,
@@ -47,6 +55,8 @@ class RobotLocalizer
 
     struct MotorData
     {
+        // Local-frame velocity, as reported directly by the motor sensors (see
+        // velocity_conversion_util.h)
         Vector velocity;
         AngularVelocity angular_velocity;
     };
@@ -113,9 +123,20 @@ class RobotLocalizer
     /**
      * Gets the estimated velocity of the robot in world space.
      *
+     * The filter estimates velocity in the robot's local frame (see StateIndex), so
+     * this converts it to world space using the current orientation estimate.
+     *
      * @return the estimated velocity of the robot in world space
      */
-    Vector getVelocity() const;
+    Vector getGlobalVelocity() const;
+
+    /**
+     * Gets the estimated velocity of the robot in its own local frame (see StateIndex
+     * and velocity_conversion_util.h), i.e. the filter's raw velocity state.
+     *
+     * @return the estimated velocity of the robot in its local frame
+     */
+    Vector getLocalVelocity() const;
 
     /**
      * Gets the estimated orientation of the robot in world space.
@@ -148,9 +169,15 @@ class RobotLocalizer
     void updateFilterWithVision(const Point& position, const Angle& orientation);
 
     /**
-     * Computes the process model, process covariance, and control model for the
-     * given elapsed time, and writes them into the filter. Does not run the
-     * predict step itself.
+     * Computes the process model function, its Jacobian, the process covariance, and
+     * the control model for the given elapsed time, and writes them into the filter.
+     * Does not run the predict step itself.
+     *
+     * The control model also depends on the filter's current orientation estimate
+     * (used to rotate the global-frame control input into the local frame that
+     * velocity is estimated in), so this must be called with the filter's state
+     * estimate set to what it was immediately before the predict step being
+     * (re)computed.
      *
      * @param delta_time_seconds The elapsed time to generate the prediction
      * matrices for
@@ -176,8 +203,9 @@ class RobotLocalizer
     {
         FilterStepType type;
 
-        // Set iff type == PREDICT. process_model/process_covariance/control_model are
-        // recomputed from the elapsed time during replay instead of being stored (see
+        // Set iff type == PREDICT. The process model function/Jacobian, process
+        // covariance, and control model are recomputed from the elapsed time and the
+        // state estimate during replay instead of being stored (see
         // generatedPredictionMatrices).
         std::optional<Eigen::Vector<double, CONTROL_SIZE>> control_input;
 
@@ -192,7 +220,7 @@ class RobotLocalizer
         double time_seconds;
     };
 
-    KalmanFilter<STATE_SIZE, MEASUREMENT_SIZE, CONTROL_SIZE> filter_;
+    ExtendedKalmanFilter<STATE_SIZE, MEASUREMENT_SIZE, CONTROL_SIZE> filter_;
 
     // Process noise variance used in prediction. The linear term models how much
     // actual velocity deviates from the commanded target velocity (a rate, per unit