Add balance control to foot controller

module/actuation/FootController/data/config/FootController.yaml

@@ -21,3 +21,18 @@ servo_gains:
   R_KNEE: 8
   R_ANKLE_PITCH: 8
   R_ANKLE_ROLL: 8
+
+balance:
+  roll_p_gain: 0.2
+  pitch_p_gain: 0.2
+
+  roll_i_gain: 0.0
+  pitch_i_gain: 0.0
+
+  roll_d_gain: 0.01
+  pitch_d_gain: 0.01
+
+  max_i_error: 0.0
+
+  max_pitch_error: 0.6
+  max_roll_error: 0.6

module/actuation/FootController/data/config/webots/FootController.yaml

@@ -1,3 +1,5 @@
+log_level: INFO
+
 startup:
   # Gain to use in startup phase
   servo_gain: 30
@@ -17,3 +19,18 @@ servo_gains:
   R_KNEE: 30
   R_ANKLE_PITCH: 30
   R_ANKLE_ROLL: 30
+
+balance:
+  roll_p_gain: 0.2
+  pitch_p_gain: 0.2
+
+  roll_i_gain: 0.0
+  pitch_i_gain: 0.0
+
+  roll_d_gain: 0.01
+  pitch_d_gain: 0.01
+
+  max_i_error: 0.0
+
+  max_pitch_error: 0.6
+  max_roll_error: 0.6

module/actuation/FootController/src/FootController.cpp

@@ -45,17 +45,29 @@ namespace module::actuation {
 
         on<Configuration>("FootController.yaml").then([this](const Configuration& config) {
             // Set log level from the configuration
-            this->log_level = config["log_level"].as<NUClear::LogLevel>();
-            cfg.mode        = config["mode"].as<std::string>();
-
+            this->log_level   = config["log_level"].as<NUClear::LogLevel>();
+            cfg.mode          = config["mode"].as<std::string>();
             cfg.desired_gains = config["servo_gains"].as<std::map<std::string, double>>();
+
             // Set gains of servo to startup phase values
             cfg.servo_states.clear();
             cfg.startup_gain = config["startup"]["servo_gain"].as<double>();
             for (const auto& servo : cfg.desired_gains) {
                 utility::input::ServoID servo_id(servo.first);
                 cfg.servo_states[servo_id] = ServoState(cfg.startup_gain, TORQUE_ENABLED);
             }
+
+            // Balance config
+            cfg.roll_p_gain     = config["balance"]["roll_p_gain"].as<double>();
+            cfg.pitch_p_gain    = config["balance"]["pitch_p_gain"].as<double>();
+            cfg.roll_i_gain     = config["balance"]["roll_i_gain"].as<double>();
+            cfg.pitch_i_gain    = config["balance"]["pitch_i_gain"].as<double>();
+            cfg.max_i_error     = config["balance"]["max_i_error"].as<double>();
+            cfg.roll_d_gain     = config["balance"]["roll_d_gain"].as<double>();
+            cfg.pitch_d_gain    = config["balance"]["pitch_d_gain"].as<double>();
+            cfg.max_pitch_error = config["balance"]["max_pitch_error"].as<double>();
+            cfg.max_roll_error  = config["balance"]["max_roll_error"].as<double>();
+
             // Emit request to set desired gains after a delay
             emit<Scope::DELAY>(std::make_unique<SetGains>(),
                                std::chrono::seconds(config["startup"]["duration"].as<int>()));
@@ -68,7 +80,6 @@ namespace module::actuation {
             }
         });
 
-
         on<Provide<ControlLeftFoot>, With<Sensors>, Needs<LeftLegIK>, Priority::HIGH>().then(
             [this](const ControlLeftFoot& left_foot, const Sensors& sensors) {
                 auto left_leg = std::make_unique<LeftLegIK>();

module/actuation/FootController/src/FootController.hpp

@@ -27,7 +27,10 @@
 #ifndef MODULE_ACTUATION_FOOTCONTROLLER_HPP
 #define MODULE_ACTUATION_FOOTCONTROLLER_HPP
 
+#include <Eigen/Core>
+#include <Eigen/Geometry>
 #include <algorithm>
+#include <fmt/format.h>
 #include <nuclear>
 
 #include "extension/Behaviour.hpp"
@@ -37,6 +40,8 @@
 
 #include "utility/input/LimbID.hpp"
 #include "utility/input/ServoID.hpp"
+#include "utility/math/comparison.hpp"
+#include "utility/math/euler.hpp"
 #include "utility/nusight/NUhelpers.hpp"
 
 #define TORQUE_ENABLED 100
@@ -47,11 +52,18 @@ namespace module::actuation {
     using message::input::Sensors;
 
     using utility::input::LimbID;
+    using utility::math::euler::mat_to_rpy_intrinsic;
+    using utility::math::euler::rpy_intrinsic_to_mat;
     using utility::nusight::graph;
 
+    /// @brief Empty struct used to trigger the setting of servo gains.
     struct SetGains {};
 
     class FootController : public ::extension::behaviour::BehaviourReactor {
+    public:
+        /// @brief Called by the powerplant to build and setup the FootController reactor.
+        explicit FootController(std::unique_ptr<NUClear::Environment> environment);
+
     private:
         /// @brief Stores configuration values
         struct Config {
@@ -79,8 +91,36 @@ namespace module::actuation {
             std::chrono::seconds set_gain_delay = std::chrono::seconds(0);
             /// @brief Desired gains for each servo
             std::map<std::string, double> desired_gains{};
+
+            /// @brief Proportional gain for torso orientation roll correction
+            double roll_p_gain = 0.0;
+            /// @brief Proportional gain for torso orientation pitch correction
+            double pitch_p_gain = 0.0;
+            /// @brief Integral gain for torso orientation roll correction
+            double roll_i_gain = 0.0;
+            /// @brief Integral gain for torso orientation pitch correction
+            double pitch_i_gain = 0.0;
+            /// @brief Antiwindup max I error
+            double max_i_error = 0.0;
+            /// @brief Derivative gain for torso orientation roll correction
+            double roll_d_gain = 0.0;
+            /// @brief Derivative gain for torso orientation pitch correction
+            double pitch_d_gain = 0.0;
+            /// @brief Maximum allowed pitch error before disabling correction
+            double max_pitch_error = 0.0;
+            /// @brief Maximum allowed roll error before disabling correction
+            double max_roll_error = 0.0;
         } cfg;
 
+        /// @brief Accumulates the integral of the roll error over time for use in PID control.
+        double integral_roll_error = 0;
+        /// @brief Accumulates the integral of the pitch error over time for use in PID control.
+        double integral_pitch_error = 0;
+
+        /// @brief Stores the timestamp of the last PID control update, used to compute time deltas for integral and
+        /// derivative calculations.
+        NUClear::clock::time_point last_update_time{};
+
         /// @brief Foot controller logic
         template <typename FootControlTask, typename IKTask>
         void control_foot(const FootControlTask& foot_control_task,
@@ -110,13 +150,127 @@ namespace module::actuation {
             else {
                 throw std::runtime_error("Invalid mode");
             }
-        }
 
-    public:
-        /// @brief Called by the powerplant to build and setup the FootController reactor.
-        explicit FootController(std::unique_ptr<NUClear::Environment> environment);
+            // Add correction to the torso orientation if enabled
+            if (foot_control_task.correction_enabled) {
+                // Hwt quaternion
+                Eigen::Quaterniond Hwt_quat(sensors.Htw.inverse().linear());
+
+                // Get fused roll and pitch
+                Eigen::Vector3d rpy = mat_to_rpy_intrinsic(Hwt_quat.toRotationMatrix());
+                double fused_roll   = rpy.x();
+                double fused_pitch  = rpy.y();
+
+                if (log_level <= DEBUG) {
+                    // Graph the correction being applied
+                    emit(graph("Balance/Actual_Roll", fused_roll));
+                    emit(graph("Balance/Actual_Pitch", fused_pitch));
+                }
+
+                // Get the desired roll and pitch
+                Eigen::Quaterniond Hft_quat(ik_task->Htf.inverse().linear());
+                rpy                  = mat_to_rpy_intrinsic(Hft_quat.toRotationMatrix());
+                double desired_roll  = rpy.x();
+                double desired_pitch = rpy.y();
+
+                if (log_level <= DEBUG) {
+                    // Graph desired orientation (before PID correction)
+                    emit(graph("Balance/Desired_Roll_Original", desired_roll));
+                    emit(graph("Balance/Desired_Pitch_Original", desired_pitch));
+                }
+
+                // Compute the error between the desired torso orientation and the actual torso orientation
+                auto roll_error  = desired_roll - fused_roll;
+                auto pitch_error = desired_pitch - fused_pitch;
+
+                if (log_level <= DEBUG) {
+                    // Graph the errors
+                    emit(graph("Balance/Roll_Error", roll_error));
+                    emit(graph("Balance/Pitch_Error", pitch_error));
+                }
+
+                auto dt =
+                    std::chrono::duration_cast<std::chrono::duration<double>>(NUClear::clock::now() - last_update_time)
+                        .count();
+                last_update_time = NUClear::clock::now();
+
+                // P control
+                auto clamped_roll_error  = utility::math::clamp(-cfg.max_roll_error, roll_error, cfg.max_roll_error);
+                auto clamped_pitch_error = utility::math::clamp(-cfg.max_pitch_error, pitch_error, cfg.max_pitch_error);
+                desired_roll += cfg.roll_p_gain * clamped_roll_error;
+                desired_pitch += cfg.pitch_p_gain * clamped_pitch_error;
+
+                if (log_level <= DEBUG) {
+                    if (roll_error > cfg.max_roll_error || roll_error < -cfg.max_roll_error) {
+                        // If the error is too large, we are probably falling over and should pause applying control.
+                        log<DEBUG>(fmt::format("Balance correction disabled due to large ROLL error: roll_error = {}",
+                                               roll_error));
+                        log<DEBUG>(fmt::format("Clamping roll error to: {}", clamped_roll_error));
+                    }
+
+                    if (pitch_error > cfg.max_pitch_error || pitch_error < -cfg.max_pitch_error) {
+                        // If the error is too large, we are probably falling over and should pause applying control.
+                        log<DEBUG>(fmt::format("Balance correction disabled due to large error: pitch_error = {}",
+                                               pitch_error));
+                        log<DEBUG>(fmt::format("Clamping pitch error to: {}", clamped_pitch_error));
+                    }
+                }
+
+                // Graph P Correction
+                if (log_level <= DEBUG) {
+                    emit(graph("Balance/P_Roll_Correction", cfg.roll_p_gain * roll_error));
+                    emit(graph("Balance/P_Pitch_Correction", cfg.pitch_p_gain * pitch_error));
+                }
+
+                // I control
+                integral_roll_error += roll_error * dt;
+                integral_pitch_error += pitch_error * dt;
+
+                // Graph the integral correction
+                if (log_level <= DEBUG) {
+                    emit(graph("Balance/I_Roll_Correction", cfg.roll_i_gain * integral_roll_error));
+                    emit(graph("Balance/I_Pitch_Correction", cfg.pitch_i_gain * integral_pitch_error));
+                }
+
+                // Anti windup
+                integral_roll_error  = std::max(std::min(integral_roll_error, cfg.max_i_error), -cfg.max_i_error);
+                integral_pitch_error = std::max(std::min(integral_pitch_error, cfg.max_i_error), -cfg.max_i_error);
+
+                desired_roll += cfg.roll_i_gain * integral_roll_error;
+                desired_pitch += cfg.pitch_i_gain * integral_pitch_error;
+
+                // D control using gyroscope values
+                auto roll_rate  = sensors.gyroscope.x();
+                auto pitch_rate = sensors.gyroscope.y();
+                desired_roll -= cfg.roll_d_gain * roll_rate;
+                desired_pitch -= cfg.pitch_d_gain * pitch_rate;
+
+                if (log_level <= DEBUG) {
+                    // Graph D corrections and rates
+                    emit(graph("Balance/D_Roll_Correction", -cfg.roll_d_gain * roll_rate));
+                    emit(graph("Balance/D_Pitch_Correction", -cfg.pitch_d_gain * pitch_rate));
+                    emit(graph("Balance/Roll_Rate", roll_rate));
+                    emit(graph("Balance/Pitch_Rate", pitch_rate));
+
+                    // Graph final corrected desired values
+                    emit(graph("Balance/Desired_Roll_Final", desired_roll));
+                    emit(graph("Balance/Desired_Pitch_Final", desired_pitch));
+                }
+
+                double desired_yaw = mat_to_rpy_intrinsic(Hft_quat.toRotationMatrix()).z();
+
+                // Compute desired orientation: yaw * fused_roll_pitch
+                Eigen::Matrix3d desired_Rft = Eigen::AngleAxisd(desired_yaw, Eigen::Vector3d::UnitZ())
+                                              * rpy_intrinsic_to_mat(Eigen::Vector3d(desired_roll, desired_pitch, 0.0));
+                Eigen::Isometry3d Htf_corrected = foot_control_task.Htf;
+                Htf_corrected.linear()          = desired_Rft.transpose();
+
+                ik_task->Htf = Htf_corrected;
+            }
+        }
     };
 
+
 }  // namespace module::actuation
 
 #endif  // MODULE_ACTUATION_FOOTCONTROLLER_HPP

module/skill/SplineKick/data/config/SplineKick.yaml

@@ -37,3 +37,6 @@ arms:
   left_shoulder_roll: 0.35
   right_elbow: -0.7
   left_elbow: -0.7
+
+# Use balance control. NOTE: correction_enabled must be true in FootController config
+use_balance_control: false

module/skill/SplineKick/data/config/webots/SplineKick.yaml

@@ -36,3 +36,6 @@ arms:
   left_shoulder_roll: 0.2
   right_elbow: -0.7
   left_elbow: -0.7
+
+# Use balance control. NOTE: correction_enabled must be true in FootController config
+use_balance_control: false

module/skill/SplineKick/src/SplineKick.cpp

@@ -98,6 +98,8 @@ namespace module::skill {
             cfg.arm_positions.emplace_back(ServoID::L_SHOULDER_ROLL, config["arms"]["left_shoulder_roll"].as<double>());
             cfg.arm_positions.emplace_back(ServoID::R_ELBOW, config["arms"]["right_elbow"].as<double>());
             cfg.arm_positions.emplace_back(ServoID::L_ELBOW, config["arms"]["left_elbow"].as<double>());
+            // Configure balance control
+            cfg.use_balance_control = config["use_balance_control"].as<bool>();
         });
 
         // Start - Runs every time the Kick provider starts (wasn't running)
@@ -139,8 +141,8 @@ namespace module::skill {
                 Eigen::Isometry3d Htr = kick_generator.get_foot_pose(LimbID::RIGHT_LEG);
 
                 // Construct ControlFoot tasks
-                emit<Task>(std::make_unique<ControlLeftFoot>(Htl, goal_time));
-                emit<Task>(std::make_unique<ControlRightFoot>(Htr, goal_time));
+                emit<Task>(std::make_unique<ControlLeftFoot>(Htl, goal_time, cfg.use_balance_control));
+                emit<Task>(std::make_unique<ControlRightFoot>(Htr, goal_time, cfg.use_balance_control));
 
                 // Construct Arm IK tasks
                 auto left_arm  = std::make_unique<LeftArm>();

module/skill/SplineKick/src/SplineKick.hpp

@@ -53,6 +53,9 @@ namespace module::skill {
 
             /// @brief Desired arm positions while walking
             std::vector<std::pair<utility::input::ServoID, double>> arm_positions{};
+
+            /// @brief Option to use balance control when kicking
+            bool use_balance_control = false;
         } cfg;
 
         /// @brief Last time we updated the walk engine

module/skill/Walk/data/config/Walk.yaml

@@ -5,6 +5,9 @@ walk:
   period: 0.39
   # Bool to indicate whether or not to only switch if the next foot is planted
   only_switch_when_planted: true
+  # Bool to indicate whether or not to use balance control. NOTE: correction_enabled must be true in FootController config
+  use_balance_control: true
+
   step:
     # Maximum step limits in x and y (in meters), and theta (in radians).
     limits: [0.5, 0.2, 0.4]

module/skill/Walk/data/config/webots/Walk.yaml

@@ -5,6 +5,9 @@ walk:
   period: 0.39
   # Bool to indicate whether or not to only switch if the next foot is planted
   only_switch_when_planted: true
+  # Bool to indicate whether or not to use balance control. NOTE: correction_enabled must be true in FootController config
+  use_balance_control: false
+
   step:
     # Maximum step limits in x and y (in meters), and theta (in radians).
     limits: [0.5, 0.2, 0.4]

module/skill/Walk/src/Walk.cpp

@@ -107,9 +107,10 @@ namespace module::skill {
             cfg.arm_positions.emplace_back(ServoID::L_SHOULDER_ROLL, config["arms"]["left_shoulder_roll"].as<double>());
             cfg.arm_positions.emplace_back(ServoID::R_ELBOW, config["arms"]["right_elbow"].as<double>());
             cfg.arm_positions.emplace_back(ServoID::L_ELBOW, config["arms"]["left_elbow"].as<double>());
-            cfg.kick_velocity_x    = config["kick"]["kick_velocity_x"].as<double>();
-            cfg.kick_velocity_y    = config["kick"]["kick_velocity_y"].as<double>();
-            cfg.kick_timing_offset = config["kick"]["kick_timing_offset"].as<double>();
+            cfg.kick_velocity_x     = config["kick"]["kick_velocity_x"].as<double>();
+            cfg.kick_velocity_y     = config["kick"]["kick_velocity_y"].as<double>();
+            cfg.kick_timing_offset  = config["kick"]["kick_timing_offset"].as<double>();
+            cfg.use_balance_control = config["walk"]["use_balance_control"].as<bool>();
 
             // Since walk needs a Stability message to run, emit one at the beginning
             emit(std::make_unique<Stability>(Stability::UNKNOWN));
@@ -212,8 +213,8 @@ namespace module::skill {
                 Eigen::Isometry3d Htr = walk_generator.get_foot_pose(LimbID::RIGHT_LEG);
 
                 // Construct ControlFoot tasks
-                emit<Task>(std::make_unique<ControlLeftFoot>(Htl, goal_time));
-                emit<Task>(std::make_unique<ControlRightFoot>(Htr, goal_time));
+                emit<Task>(std::make_unique<ControlLeftFoot>(Htl, goal_time, cfg.use_balance_control));
+                emit<Task>(std::make_unique<ControlRightFoot>(Htr, goal_time, cfg.use_balance_control));
 
                 // Construct Arm IK tasks
                 auto left_arm  = std::make_unique<LeftArm>();

module/skill/Walk/src/Walk.hpp

@@ -66,6 +66,8 @@ namespace module::skill {
             double kick_velocity_y = 0.0;
             /// @brief Buffer time before the kick starts
             double kick_timing_offset = 0.0;
+            /// @brief Option to use balance control when walking
+            bool use_balance_control = false;
         } cfg;
 
         /// @brief Last time we updated the walk engine