Add roll joint to superstructure and arm UI

Arm UI changes:
- Update robot dimensions
- Support visualizing roll joint
- Add roll joint collision detection

Superstructure changes:
- Adding roll joint feedback loop and zeroing estimator

Signed-off-by: milind-u <milind.upadhyay@gmail.com>
Change-Id: I422e343890248940bba98ba3cabac94e68723a3e
diff --git a/y2023/control_loops/superstructure/arm/BUILD b/y2023/control_loops/superstructure/arm/BUILD
index 0ef2cf9..032c654 100644
--- a/y2023/control_loops/superstructure/arm/BUILD
+++ b/y2023/control_loops/superstructure/arm/BUILD
@@ -10,15 +10,15 @@
     visibility = ["//visibility:public"],
     deps = [
         ":generated_graph",
-        "//frc971/control_loops/double_jointed_arm:demo_path",
         "//frc971/control_loops/double_jointed_arm:ekf",
         "//frc971/control_loops/double_jointed_arm:graph",
-        "//frc971/control_loops/double_jointed_arm:trajectory",
         "//frc971/zeroing",
         "//y2023:constants",
         "//y2023/control_loops/superstructure:superstructure_position_fbs",
         "//y2023/control_loops/superstructure:superstructure_status_fbs",
         "//y2023/control_loops/superstructure/arm:arm_constants",
+        "//y2023/control_loops/superstructure/arm:trajectory",
+        "//y2023/control_loops/superstructure/roll:roll_plants",
     ],
 )
 
@@ -49,7 +49,8 @@
     deps = [
         ":arm_constants",
         "//frc971/control_loops/double_jointed_arm:graph",
-        "//frc971/control_loops/double_jointed_arm:trajectory",
+        "//y2023/control_loops/superstructure/arm:trajectory",
+        "//y2023/control_loops/superstructure/roll:roll_plants",
     ],
 )
 
diff --git a/y2023/control_loops/superstructure/arm/arm.cc b/y2023/control_loops/superstructure/arm/arm.cc
index 73781d2..5417dc8 100644
--- a/y2023/control_loops/superstructure/arm/arm.cc
+++ b/y2023/control_loops/superstructure/arm/arm.cc
@@ -1,5 +1,8 @@
 #include "y2023/control_loops/superstructure/arm/arm.h"
 
+#include "y2023/control_loops/superstructure/roll/integral_hybrid_roll_plant.h"
+#include "y2023/control_loops/superstructure/roll/integral_roll_plant.h"
+
 namespace y2023 {
 namespace control_loops {
 namespace superstructure {
@@ -18,20 +21,24 @@
       state_(ArmState::UNINITIALIZED),
       proximal_zeroing_estimator_(values_->arm_proximal.zeroing),
       distal_zeroing_estimator_(values_->arm_distal.zeroing),
+      roll_joint_zeroing_estimator_(values_->roll_joint.zeroing),
       proximal_offset_(0.0),
       distal_offset_(0.0),
-      max_intake_override_(1000.0),
-      alpha_unitizer_((::Eigen::Matrix<double, 2, 2>() << 1.0 / kAlpha0Max(),
-                       0.0, 0.0, 1.0 / kAlpha1Max())
+      roll_joint_offset_(0.0),
+      alpha_unitizer_((::Eigen::DiagonalMatrix<double, 3>().diagonal()
+                           << (1.0 / kAlpha0Max()),
+                       (1.0 / kAlpha1Max()), (1.0 / kAlpha2Max()))
                           .finished()),
       dynamics_(kArmConstants),
-      search_graph_(MakeSearchGraph(&dynamics_, &trajectories_, alpha_unitizer_,
-                                    kVMax())),
       close_enough_for_full_power_(false),
       brownout_count_(0),
+      roll_joint_loop_(roll::MakeIntegralRollLoop()),
+      hybrid_roll_joint_loop_(roll::MakeIntegralHybridRollLoop()),
       arm_ekf_(&dynamics_),
+      search_graph_(MakeSearchGraph(&dynamics_, &trajectories_, alpha_unitizer_,
+                                    kVMax(), &hybrid_roll_joint_loop_)),
       // Go to the start of the first trajectory.
-      follower_(&dynamics_, NeutralPosPoint()),
+      follower_(&dynamics_, &hybrid_roll_joint_loop_, NeutralPosPoint()),
       points_(PointList()),
       current_node_(0) {
   int i = 0;
@@ -48,26 +55,32 @@
     const ::aos::monotonic_clock::time_point /*monotonic_now*/,
     const uint32_t *unsafe_goal, const superstructure::ArmPosition *position,
     bool trajectory_override, double *proximal_output, double *distal_output,
-    bool /*intake*/, bool /*spit*/, flatbuffers::FlatBufferBuilder *fbb) {
-  ::Eigen::Matrix<double, 2, 1> Y;
+    double *roll_joint_output, bool /*intake*/, bool /*spit*/,
+    flatbuffers::FlatBufferBuilder *fbb) {
   const bool outputs_disabled =
-      ((proximal_output == nullptr) || (distal_output == nullptr));
+      ((proximal_output == nullptr) || (distal_output == nullptr) ||
+       (roll_joint_output == nullptr));
   if (outputs_disabled) {
     ++brownout_count_;
   } else {
     brownout_count_ = 0;
   }
 
-  uint32_t filtered_goal = 0;
+  // TODO(milind): should we default to the closest position?
+  uint32_t filtered_goal = arm::NeutralPosIndex();
   if (unsafe_goal != nullptr) {
     filtered_goal = *unsafe_goal;
   }
 
-  Y << position->proximal()->encoder() + proximal_offset_,
+  ::Eigen::Matrix<double, 2, 1> Y_arm;
+  Y_arm << position->proximal()->encoder() + proximal_offset_,
       position->distal()->encoder() + distal_offset_;
+  ::Eigen::Matrix<double, 1, 1> Y_roll_joint;
+  Y_roll_joint << position->roll_joint()->encoder() + roll_joint_offset_;
 
   proximal_zeroing_estimator_.UpdateEstimate(*position->proximal());
   distal_zeroing_estimator_.UpdateEstimate(*position->distal());
+  roll_joint_zeroing_estimator_.UpdateEstimate(*position->roll_joint());
 
   if (proximal_output != nullptr) {
     *proximal_output = 0.0;
@@ -75,15 +88,21 @@
   if (distal_output != nullptr) {
     *distal_output = 0.0;
   }
+  if (roll_joint_output != nullptr) {
+    *roll_joint_output = 0.0;
+  }
 
-  arm_ekf_.Correct(Y, kDt());
+  arm_ekf_.Correct(Y_arm, kDt());
+  roll_joint_loop_.Correct(Y_roll_joint);
 
   if (::std::abs(arm_ekf_.X_hat(0) - follower_.theta(0)) <= 0.05 &&
-      ::std::abs(arm_ekf_.X_hat(2) - follower_.theta(1)) <= 0.05) {
+      ::std::abs(arm_ekf_.X_hat(2) - follower_.theta(1)) <= 0.05 &&
+      ::std::abs(roll_joint_loop_.X_hat(0) - follower_.theta(2)) <= 0.05) {
     close_enough_for_full_power_ = true;
   }
   if (::std::abs(arm_ekf_.X_hat(0) - follower_.theta(0)) >= 1.10 ||
-      ::std::abs(arm_ekf_.X_hat(2) - follower_.theta(1)) >= 1.10) {
+      ::std::abs(arm_ekf_.X_hat(2) - follower_.theta(1)) >= 1.10 ||
+      ::std::abs(roll_joint_loop_.X_hat(0) - follower_.theta(2)) >= 0.50) {
     close_enough_for_full_power_ = false;
   }
 
@@ -94,25 +113,31 @@
       state_ = ArmState::ZEROING;
       proximal_zeroing_estimator_.Reset();
       distal_zeroing_estimator_.Reset();
+      roll_joint_zeroing_estimator_.Reset();
       break;
 
     case ArmState::ZEROING:
       // Zero by not moving.
-      if (proximal_zeroing_estimator_.zeroed() &&
-          distal_zeroing_estimator_.zeroed()) {
+      if (zeroed()) {
         state_ = ArmState::DISABLED;
 
         proximal_offset_ = proximal_zeroing_estimator_.offset();
         distal_offset_ = distal_zeroing_estimator_.offset();
+        roll_joint_offset_ = roll_joint_zeroing_estimator_.offset();
 
-        Y << position->proximal()->encoder() + proximal_offset_,
+        Y_arm << position->proximal()->encoder() + proximal_offset_,
             position->distal()->encoder() + distal_offset_;
+        Y_roll_joint << position->roll_joint()->encoder() + roll_joint_offset_;
 
         // TODO(austin): Offset ekf rather than reset it.  Since we aren't
         // moving at this point, it's pretty safe to do this.
-        ::Eigen::Matrix<double, 4, 1> X;
-        X << Y(0), 0.0, Y(1), 0.0;
-        arm_ekf_.Reset(X);
+        ::Eigen::Matrix<double, 4, 1> X_arm;
+        X_arm << Y_arm(0), 0.0, Y_arm(1), 0.0;
+        arm_ekf_.Reset(X_arm);
+
+        ::Eigen::Matrix<double, 3, 1> X_roll_joint;
+        X_roll_joint << Y_roll_joint(0), 0.0, 0.0;
+        roll_joint_loop_.mutable_X_hat() = X_roll_joint;
       } else {
         break;
       }
@@ -122,14 +147,14 @@
       follower_.SwitchTrajectory(nullptr);
       close_enough_for_full_power_ = false;
 
-      const ::Eigen::Matrix<double, 2, 1> current_theta =
-          (::Eigen::Matrix<double, 2, 1>() << arm_ekf_.X_hat(0),
-           arm_ekf_.X_hat(2))
+      const ::Eigen::Matrix<double, 3, 1> current_theta =
+          (::Eigen::Matrix<double, 3, 1>() << arm_ekf_.X_hat(0),
+           arm_ekf_.X_hat(2), roll_joint_loop_.X_hat(0))
               .finished();
       uint32_t best_index = 0;
       double best_distance = (points_[0] - current_theta).norm();
       uint32_t current_index = 0;
-      for (const ::Eigen::Matrix<double, 2, 1> &point : points_) {
+      for (const ::Eigen::Matrix<double, 3, 1> &point : points_) {
         const double new_distance = (point - current_theta).norm();
         if (new_distance < best_distance) {
           best_distance = new_distance;
@@ -165,7 +190,8 @@
       // ESTOP if we hit the hard limits.
       // TODO(austin): Pick some sane limits.
       if (proximal_zeroing_estimator_.error() ||
-          distal_zeroing_estimator_.error()) {
+          distal_zeroing_estimator_.error() ||
+          roll_joint_zeroing_estimator_.error()) {
         AOS_LOG(ERROR, "Zeroing error ESTOP\n");
         state_ = ArmState::ESTOP;
       } else if (outputs_disabled && brownout_count_ > kMaxBrownoutCount) {
@@ -227,38 +253,56 @@
       close_enough_for_full_power_
           ? kOperatingVoltage()
           : (state_ == ArmState::GOTO_PATH ? kGotoPathVMax() : kPathlessVMax());
-  follower_.Update(arm_ekf_.X_hat(), disable, kDt(), vmax_,
-                   max_operating_voltage);
+  ::Eigen::Matrix<double, 9, 1> X_hat;
+  X_hat.block<6, 1>(0, 0) = arm_ekf_.X_hat();
+  X_hat.block<3, 1>(6, 0) = roll_joint_loop_.X_hat();
+
+  follower_.Update(X_hat, disable, kDt(), vmax_, max_operating_voltage);
   AOS_LOG(INFO, "Max voltage: %f\n", max_operating_voltage);
 
+  arm_ekf_.Predict(follower_.U().head<2>(), kDt());
+  roll_joint_loop_.UpdateObserver(follower_.U().tail<1>(), kDtDuration());
+
   flatbuffers::Offset<frc971::PotAndAbsoluteEncoderEstimatorState>
       proximal_estimator_state_offset =
           proximal_zeroing_estimator_.GetEstimatorState(fbb);
   flatbuffers::Offset<frc971::PotAndAbsoluteEncoderEstimatorState>
       distal_estimator_state_offset =
           distal_zeroing_estimator_.GetEstimatorState(fbb);
+  flatbuffers::Offset<frc971::PotAndAbsoluteEncoderEstimatorState>
+      roll_joint_estimator_state_offset =
+          roll_joint_zeroing_estimator_.GetEstimatorState(fbb);
 
   superstructure::ArmStatus::Builder status_builder(*fbb);
   status_builder.add_proximal_estimator_state(proximal_estimator_state_offset);
   status_builder.add_distal_estimator_state(distal_estimator_state_offset);
+  status_builder.add_roll_joint_estimator_state(
+      roll_joint_estimator_state_offset);
 
   status_builder.add_goal_theta0(follower_.theta(0));
   status_builder.add_goal_theta1(follower_.theta(1));
+  status_builder.add_goal_theta2(follower_.theta(2));
   status_builder.add_goal_omega0(follower_.omega(0));
   status_builder.add_goal_omega1(follower_.omega(1));
+  status_builder.add_goal_omega2(follower_.omega(2));
 
   status_builder.add_theta0(arm_ekf_.X_hat(0));
   status_builder.add_theta1(arm_ekf_.X_hat(2));
+  status_builder.add_theta2(roll_joint_loop_.X_hat(0));
   status_builder.add_omega0(arm_ekf_.X_hat(1));
   status_builder.add_omega1(arm_ekf_.X_hat(3));
+  status_builder.add_omega2(roll_joint_loop_.X_hat(1));
   status_builder.add_voltage_error0(arm_ekf_.X_hat(4));
   status_builder.add_voltage_error1(arm_ekf_.X_hat(5));
+  status_builder.add_voltage_error2(roll_joint_loop_.X_hat(2));
 
   if (!disable) {
     *proximal_output = ::std::max(
         -kOperatingVoltage(), ::std::min(kOperatingVoltage(), follower_.U(0)));
     *distal_output = ::std::max(
         -kOperatingVoltage(), ::std::min(kOperatingVoltage(), follower_.U(1)));
+    *roll_joint_output = ::std::max(
+        -kOperatingVoltage(), ::std::min(kOperatingVoltage(), follower_.U(2)));
   }
 
   status_builder.add_path_distance_to_go(follower_.path_distance_to_go());
@@ -269,7 +313,6 @@
   status_builder.add_state(state_);
   status_builder.add_failed_solutions(follower_.failed_solutions());
 
-  arm_ekf_.Predict(follower_.U(), kDt());
   return status_builder.Finish();
 }
 
diff --git a/y2023/control_loops/superstructure/arm/arm.h b/y2023/control_loops/superstructure/arm/arm.h
index 27588f1..9cda7fc 100644
--- a/y2023/control_loops/superstructure/arm/arm.h
+++ b/y2023/control_loops/superstructure/arm/arm.h
@@ -5,15 +5,14 @@
 #include "frc971/control_loops/double_jointed_arm/dynamics.h"
 #include "frc971/control_loops/double_jointed_arm/ekf.h"
 #include "frc971/control_loops/double_jointed_arm/graph.h"
-#include "frc971/control_loops/double_jointed_arm/trajectory.h"
 #include "frc971/zeroing/zeroing.h"
 #include "y2023/constants.h"
 #include "y2023/control_loops/superstructure/arm/generated_graph.h"
+#include "y2023/control_loops/superstructure/arm/trajectory.h"
 #include "y2023/control_loops/superstructure/superstructure_position_generated.h"
 #include "y2023/control_loops/superstructure/superstructure_status_generated.h"
 
 using frc971::control_loops::arm::EKF;
-using frc971::control_loops::arm::TrajectoryFollower;
 
 namespace y2023 {
 namespace control_loops {
@@ -32,8 +31,13 @@
     return kGrannyMode() ? 5.0 : 12.0;
   }
   static constexpr double kDt() { return 0.00505; }
+  static constexpr std::chrono::nanoseconds kDtDuration() {
+    return std::chrono::duration_cast<std::chrono::nanoseconds>(
+        std::chrono::duration<double>(kDt()));
+  }
   static constexpr double kAlpha0Max() { return kGrannyMode() ? 5.0 : 15.0; }
   static constexpr double kAlpha1Max() { return kGrannyMode() ? 5.0 : 15.0; }
+  static constexpr double kAlpha2Max() { return kGrannyMode() ? 5.0 : 15.0; }
 
   static constexpr double kVMax() { return kGrannyMode() ? 5.0 : 11.5; }
   static constexpr double kPathlessVMax() { return 5.0; }
@@ -43,7 +47,8 @@
       const ::aos::monotonic_clock::time_point /*monotonic_now*/,
       const uint32_t *unsafe_goal, const superstructure::ArmPosition *position,
       bool trajectory_override, double *proximal_output, double *distal_output,
-      bool /*intake*/, bool /*spit*/, flatbuffers::FlatBufferBuilder *fbb);
+      double *roll_joint_output, bool /*intake*/, bool /*spit*/,
+      flatbuffers::FlatBufferBuilder *fbb);
 
   void Reset();
 
@@ -52,13 +57,10 @@
   bool estopped() const { return state_ == ArmState::ESTOP; }
   bool zeroed() const {
     return (proximal_zeroing_estimator_.zeroed() &&
-            distal_zeroing_estimator_.zeroed());
+            distal_zeroing_estimator_.zeroed() &&
+            roll_joint_zeroing_estimator_.zeroed());
   }
 
-  // Returns the maximum position for the intake.  This is used to override the
-  // intake position to release the box when the state machine is lifting.
-  double max_intake_override() const { return max_intake_override_; }
-
   uint32_t current_node() const { return current_node_; }
 
   double path_distance_to_go() { return follower_.path_distance_to_go(); }
@@ -79,29 +81,36 @@
       proximal_zeroing_estimator_;
   ::frc971::zeroing::PotAndAbsoluteEncoderZeroingEstimator
       distal_zeroing_estimator_;
+  ::frc971::zeroing::PotAndAbsoluteEncoderZeroingEstimator
+      roll_joint_zeroing_estimator_;
 
   double proximal_offset_;
   double distal_offset_;
+  double roll_joint_offset_;
 
-  double max_intake_override_;
-
-  const ::Eigen::Matrix<double, 2, 2> alpha_unitizer_;
+  const ::Eigen::DiagonalMatrix<double, 3> alpha_unitizer_;
 
   double vmax_ = kVMax();
 
   frc971::control_loops::arm::Dynamics dynamics_;
 
   ::std::vector<TrajectoryAndParams> trajectories_;
-  SearchGraph search_graph_;
 
   bool close_enough_for_full_power_;
 
   size_t brownout_count_;
 
+  StateFeedbackLoop<3, 1, 1, double, StateFeedbackPlant<3, 1, 1>,
+                    StateFeedbackObserver<3, 1, 1>>
+      roll_joint_loop_;
+  const StateFeedbackLoop<3, 1, 1, double, StateFeedbackHybridPlant<3, 1, 1>,
+                          HybridKalman<3, 1, 1>>
+      hybrid_roll_joint_loop_;
   EKF arm_ekf_;
+  SearchGraph search_graph_;
   TrajectoryFollower follower_;
 
-  const ::std::vector<::Eigen::Matrix<double, 2, 1>> points_;
+  const ::std::vector<::Eigen::Matrix<double, 3, 1>> points_;
 
   // Start at the 0th index.
   uint32_t current_node_;