Unwrap gcs trajs with continous revolute joints to continous Euclidea…

…n space (#21137)

* Unwrap GCS trajectories with continous revolute joints into continous Euclidean space

* address comments

* remove empty traj case

* buildifier on build

* address platform comments

* fix comment

planning/trajectory_optimization/BUILD.bazel

@@ -176,6 +176,7 @@ drake_cc_googletest(
         "//common/test_utilities:eigen_matrix_compare",
         "//common/test_utilities:expect_no_throw",
         "//common/test_utilities:expect_throws_message",
+        "//common/trajectories:piecewise_polynomial",
         "//solvers:gurobi_solver",
         "//solvers:mosek_solver",
     ],

planning/trajectory_optimization/gcs_trajectory_optimization.cc

@@ -1158,13 +1158,107 @@ GcsTrajectoryOptimization::NormalizeSegmentTimes(
       start_time += 1.0;
     } else {
       throw std::runtime_error(
-          "All segments in the gcs trajectory must be of type "
+          "NormalizeSegmentTimes: All segments in the gcs trajectory "
+          "must be of type "
           "BezierCurve<double>.");
     }
   }
   return CompositeTrajectory<double>(normalized_bezier_curves);
 }
 
+namespace {
+bool IsMultipleOf2Pi(double value) {
+  // We allow some tolerance for trajectories coming out of GCS.
+  // TODO(sadra): find out what tolerance is appropriate.
+  double kTol = 1e-10;  // To match what is provided in the docs.
+  return std::abs(value - 2 * M_PI * std::round(value / (2 * M_PI))) < kTol;
+}
+
+// Unwrap the angle to the range [2π * round, 2π * (round+1)).
+double UnwrapAngle(const double angle, const int round) {
+  const int twopi_factors_to_remove =
+      static_cast<int>(std::floor(angle / (2 * M_PI))) - round;
+  return angle - 2 * M_PI * twopi_factors_to_remove;
+}
+}  // namespace
+
+trajectories::CompositeTrajectory<double>
+GcsTrajectoryOptimization::UnwrapToContinousTrajectory(
+    const trajectories::CompositeTrajectory<double>& gcs_trajectory,
+    std::vector<int> continuous_revolute_joints,
+    std::optional<std::vector<int>> starting_rounds) {
+  if (starting_rounds.has_value()) {
+    DRAKE_THROW_UNLESS(starting_rounds->size() ==
+                       continuous_revolute_joints.size());
+  }
+  // TODO(@anyone): make this a unique_ptr and use std::move to avoid copying.
+  std::vector<copyable_unique_ptr<Trajectory<double>>> unwrapped_trajectories;
+  int dim = gcs_trajectory.rows();
+  geometry::optimization::internal::ThrowsForInvalidContinuousJointsList(
+      dim, continuous_revolute_joints);
+  Eigen::VectorXd last_segment_finish;
+  for (int i = 0; i < gcs_trajectory.get_number_of_segments(); ++i) {
+    const auto& traj_segment = gcs_trajectory.segment(i);
+    const auto* bezier_segment =
+        dynamic_cast<const BezierCurve<double>*>(&traj_segment);
+    if (bezier_segment == nullptr) {
+      throw std::runtime_error(
+          "UnwrapToContinuousTrajectory: All segments in the gcs_trajectory "
+          "must be of type "
+          "BezierCurve<double>.");
+    }
+    Eigen::MatrixXd new_control_points = bezier_segment->control_points();
+    const Eigen::MatrixXd& old_control_points =
+        bezier_segment->control_points();
+    const Eigen::VectorXd& old_start = old_control_points.col(0);
+    std::vector<double> shift;
+    if (i == 0) {
+      // there is no shift from previous segment.
+      if (starting_rounds.has_value()) {
+        for (int j = 0; j < ssize(continuous_revolute_joints); ++j) {
+          const int joint_index = continuous_revolute_joints.at(j);
+          const int start_round = starting_rounds->at(j);
+          // This value will be subtracted from the old_start to get the
+          // shift.
+          const double joint_shift =
+              old_start(joint_index) -
+              UnwrapAngle(old_start(joint_index), start_round);
+          DRAKE_DEMAND(IsMultipleOf2Pi(joint_shift));
+          shift.push_back(joint_shift);
+        }
+      } else {
+        shift = std::vector<double>(ssize(continuous_revolute_joints), 0);
+      }
+    } else {
+      DRAKE_DEMAND(last_segment_finish.rows() == gcs_trajectory.rows());
+      // See how much shift is needed to match the previous new segment.
+      for (const int joint_index : continuous_revolute_joints) {
+        const double joint_shift =
+            old_start(joint_index) - last_segment_finish(joint_index);
+        if (!IsMultipleOf2Pi(joint_shift)) {
+          throw std::runtime_error(
+              fmt::format("UnwrapToContinuousTrajectory: The shift from "
+                          "previous segment: {} is not a multiple "
+                          "of 2π at segment {}, joint {}.",
+                          joint_shift, i, joint_index));
+        }
+        shift.push_back(joint_shift);
+      }
+    }
+    for (int j = 0; j < ssize(continuous_revolute_joints); ++j) {
+      const int joint_index = continuous_revolute_joints[j];
+      // Shift all the columns of the control points by the shift.
+      new_control_points.row(joint_index) -=
+          Eigen::VectorXd::Constant(old_control_points.cols(), shift.at(j));
+    }
+    last_segment_finish = new_control_points.rightCols(1);
+    unwrapped_trajectories.emplace_back(std::make_unique<BezierCurve<double>>(
+        bezier_segment->start_time(), bezier_segment->end_time(),
+        new_control_points));
+  }
+  return CompositeTrajectory<double>(unwrapped_trajectories);
+}
+
 std::vector<int> GetContinuousRevoluteJointIndices(
     const multibody::MultibodyPlant<double>& plant) {
   std::vector<int> indices;
@@ -1181,9 +1275,9 @@ std::vector<int> GetContinuousRevoluteJointIndices(
       }
       continue;
     }
-    // The second possibility we check for is a planar joint. If it is (and the
-    // angle component has no joint limits), we only add the third entry of the
-    // position vector, corresponding to theta.
+    // The second possibility we check for is a planar joint. If it is (and
+    // the angle component has no joint limits), we only add the third entry
+    // of the position vector, corresponding to theta.
     if (joint.type_name() == "planar") {
       if (joint.position_lower_limits()[2] ==
               -std::numeric_limits<float>::infinity() &&
@@ -1193,8 +1287,8 @@ std::vector<int> GetContinuousRevoluteJointIndices(
       }
       continue;
     }
-    // TODO(cohnt): Determine if other joint types (e.g. UniversalJoint) can be
-    // handled appropriately with wraparound edges, and if so, return their
+    // TODO(cohnt): Determine if other joint types (e.g. UniversalJoint) can
+    // be handled appropriately with wraparound edges, and if so, return their
     // indices as well.
   }
   return indices;

planning/trajectory_optimization/gcs_trajectory_optimization.h

@@ -607,6 +607,44 @@ class GcsTrajectoryOptimization final {
   static trajectories::CompositeTrajectory<double> NormalizeSegmentTimes(
       const trajectories::CompositeTrajectory<double>& trajectory);
 
+  /** Unwraps a trajectory with continuous revolute joints into a continuous
+   trajectory in the Euclidean space. Trajectories produced by
+   GcsTrajectoryOptimization for robotic systems with continuous revolute joints
+   may include apparent discontinuities, where a multiple of 2π is
+   instantaneously added to a joint value at the boundary between two adjacent
+   segments of the trajectory. This function removes such discontinuities by
+   adding or subtracting the appropriate multiple of 2π, "unwrapping" the
+   trajectory into a continuous representation suitable for downstream tasks
+   that do not take the joint wraparound into account.
+   @param gcs_trajectory The trajectory to unwrap.
+   @param continuous_revolute_joints The indices of the continuous revolute
+   joints.
+   @param starting_rounds A vector of integers that sets the starting rounds for
+   each continuous revolute joint. Given integer k for the starting_round of a
+   joint, its initial position will be wrapped into [2πk , 2π(k+1)). If the
+   starting rounds are not provided, the initial position of @p trajectory will
+   be unchanged.
+
+   @returns an unwrapped (continous in Euclidean space) CompositeTrajectory.
+
+   @throws std::exception if
+   starting_rounds.size()!=continuous_revolute_joints.size().
+   @throws std::exception if continuous_revolute_joints contain repeated indices
+   and/or indices outside the range [0, gcs_trajectory.rows()).
+   @throws std::exception if the gcs_trajectory is not continuous on the
+   manifold defined by the continuous_revolute_joints, i.e., the shift between
+   two consecutive segments is not an integer multiple of 2π (within a tolerance
+   of 1e-10 radians).
+   @throws std::exception if all the segments are not of type BezierCurve.
+   Other types are not supported yet. Note that currently the output of
+   GcsTrajectoryOptimization::SolvePath() is a CompositeTrajectory of
+   BezierCurves.
+    */
+  static trajectories::CompositeTrajectory<double> UnwrapToContinousTrajectory(
+      const trajectories::CompositeTrajectory<double>& gcs_trajectory,
+      std::vector<int> continuous_revolute_joints,
+      std::optional<std::vector<int>> starting_rounds = std::nullopt);
+
  private:
   const int num_positions_;
   const std::vector<int> continuous_revolute_joints_;

planning/trajectory_optimization/test/gcs_trajectory_optimization_test.cc

@@ -6,6 +6,7 @@
 #include "drake/common/test_utilities/eigen_matrix_compare.h"
 #include "drake/common/test_utilities/expect_no_throw.h"
 #include "drake/common/test_utilities/expect_throws_message.h"
+#include "drake/common/trajectories/piecewise_polynomial.h"
 #include "drake/geometry/optimization/geodesic_convexity.h"
 #include "drake/geometry/optimization/graph_of_convex_sets.h"
 #include "drake/geometry/optimization/hpolyhedron.h"
@@ -859,6 +860,105 @@ GTEST_TEST(GcsTrajectoryOptimizationTest, InvalidSubspace) {
       "Subspace must be a Point or HPolyhedron.");
 }
 
+GTEST_TEST(GcsTrajectoryOptimizationTest, UnwrapToContinousTrajectory) {
+  std::vector<copyable_unique_ptr<trajectories::Trajectory<double>>> segments;
+  Eigen::MatrixXd control_points_1(3, 3), control_points_2(3, 3),
+      control_points_3(3, 3);
+  // clang-format off
+  control_points_1 << 0, 1.0, 2.0,
+                      1.0 - 8 * M_PI, 2.0 - 8 * M_PI, 3.0 - 8 * M_PI,
+                      2.0 + 4 * M_PI, 1.0 + 4 * M_PI, 4.0 + 4 * M_PI;
+  control_points_2 << 2.0 , 2.5, 4.0,
+                      3.0 + 2 * M_PI, 1.0 + 2 * M_PI, 0.0 + 2 * M_PI,
+                      4.0 - 6 * M_PI, 3.0 - 6 * M_PI, 2.0 - 6 * M_PI;
+  control_points_3 << 4.0, -2.0, 0.0,
+                      0.0, 1.0, 2.0,
+                      2.0, 3.0, 4.0;
+  // clang-format on
+  segments.emplace_back(std::make_unique<trajectories::BezierCurve<double>>(
+      0.0, 1.0, control_points_1));
+  segments.emplace_back(std::make_unique<trajectories::BezierCurve<double>>(
+      1.0, 4.0, control_points_2));
+  segments.emplace_back(std::make_unique<trajectories::BezierCurve<double>>(
+      4.0, 6.0, control_points_3));
+  const auto traj = trajectories::CompositeTrajectory<double>(segments);
+  std::vector<int> continuous_revolute_joints = {1, 2};
+  const auto unwrapped_traj =
+      GcsTrajectoryOptimization::UnwrapToContinousTrajectory(
+          traj, continuous_revolute_joints);
+  // Small number to shift time around discontinuity points.
+  const double time_eps = 1e-7;
+  // Small number to compare position around discontinuity points.
+  // A rough upper bound for the largest derivative of the Bezier curve is about
+  // 10, therefore the error upper bound is expected to be around 10 * time_eps.
+  const double pos_eps = 1e-6;
+  double middle_time_1 = unwrapped_traj.get_segment_times()[1];
+  double middle_time_2 = unwrapped_traj.get_segment_times()[2];
+  EXPECT_NEAR(middle_time_1, 1.0, time_eps);
+  EXPECT_NEAR(middle_time_2, 4.0, time_eps);
+  // Verify the unwrapped trajectory is continous at the middle times.
+  EXPECT_TRUE(CompareMatrices(unwrapped_traj.value(middle_time_1 - time_eps),
+                              unwrapped_traj.value(middle_time_1 + time_eps),
+                              pos_eps));
+  EXPECT_TRUE(CompareMatrices(unwrapped_traj.value(middle_time_2 - time_eps),
+                              unwrapped_traj.value(middle_time_2 + time_eps),
+                              pos_eps));
+  std::vector<int> starting_rounds = {-1, 0};
+  const auto unwrapped_traj_with_start =
+      GcsTrajectoryOptimization::UnwrapToContinousTrajectory(
+          traj, continuous_revolute_joints, starting_rounds);
+  // Check if the start is unwrapped to the correct value.
+  EXPECT_TRUE(CompareMatrices(unwrapped_traj_with_start.value(0.0),
+                              Eigen::Vector3d{0.0, 1.0 - 2 * M_PI, 2.0},
+                              pos_eps));
+  EXPECT_TRUE(CompareMatrices(
+      unwrapped_traj_with_start.value(middle_time_1 - time_eps),
+      unwrapped_traj_with_start.value(middle_time_1 + time_eps), pos_eps));
+  EXPECT_TRUE(CompareMatrices(
+      unwrapped_traj_with_start.value(middle_time_2 - time_eps),
+      unwrapped_traj_with_start.value(middle_time_2 + time_eps), pos_eps));
+  // Check for invalid start_rounds
+  const std::vector<int> invalid_start_rounds = {-1, 0, 1};
+  EXPECT_THROW(GcsTrajectoryOptimization::UnwrapToContinousTrajectory(
+                   traj, continuous_revolute_joints, invalid_start_rounds),
+               std::runtime_error);
+  // Check for discontinuity for continuous revolute joints
+  control_points_2 << 2.5, 2.5, 4.0, 3.0 + 2 * M_PI, 1.0 + 2 * M_PI,
+      0.0 + 2 * M_PI, 4.0 - 6 * M_PI, 3.0 - 6 * M_PI, 2.0 - 6 * M_PI;
+  segments[1] = std::make_unique<trajectories::BezierCurve<double>>(
+      1.0, 4.0, control_points_2);
+  const auto traj_not_continous_on_revolute_manifold =
+      trajectories::CompositeTrajectory<double>(segments);
+  // For joint 0, the trajectory is not continous: 2.0 (end of segment 0)--> 2.5
+  // (start of segment 1). If we treat joint 0 as continous revolute joint, the
+  // unwrapping will fail.
+  continuous_revolute_joints = {0, 1};
+  DRAKE_EXPECT_THROWS_MESSAGE(
+      GcsTrajectoryOptimization::UnwrapToContinousTrajectory(
+          traj_not_continous_on_revolute_manifold, continuous_revolute_joints),
+      ".*is not a multiple of 2π at segment.*");
+}
+
+GTEST_TEST(GcsTrajectoryOptimizationTest, NotBezierCurveError) {
+  std::vector<copyable_unique_ptr<trajectories::Trajectory<double>>> segments;
+  auto traj_1 = trajectories::PiecewisePolynomial<double>::FirstOrderHold(
+      std::vector<double>{0, 1},
+      std::vector<Eigen::MatrixXd>{Eigen::MatrixXd::Zero(1, 2),
+                                   Eigen::MatrixXd::Ones(1, 2)});
+  auto traj_2 = trajectories::PiecewisePolynomial<double>::FirstOrderHold(
+      std::vector<double>{1, 2},
+      std::vector<Eigen::MatrixXd>{Eigen::MatrixXd::Zero(1, 2),
+                                   Eigen::MatrixXd::Ones(1, 2)});
+  segments.emplace_back(traj_1.Clone());
+  segments.emplace_back(traj_2.Clone());
+  const auto traj = trajectories::CompositeTrajectory<double>(segments);
+  std::vector<int> continuous_revolute_joints = {0};
+  DRAKE_EXPECT_THROWS_MESSAGE(
+      GcsTrajectoryOptimization::UnwrapToContinousTrajectory(
+          traj, continuous_revolute_joints),
+      ".*BezierCurve<double>.*");
+}
+
 /* This 2D environment has been presented in the GCS paper.
 ░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░
 ░░      ░░░░░░░░                                      ░░░░░░░░         ░░