Effort limits for Kinematic Trajectory Optimization

Resolves #22500.

Author: RussTedrake

bindings/pydrake/planning/BUILD.bazel

@@ -105,6 +105,9 @@ drake_py_unittest(
 
 drake_py_unittest(
     name = "trajectory_optimization_test",
+    data = [
+        "//examples/multibody/cart_pole:models",
+    ],
     deps = [
         ":planning",
         "//bindings/pydrake/examples",

bindings/pydrake/planning/planning_py_trajectory_optimization.cc

@@ -332,6 +332,11 @@ void DefinePlanningTrajectoryOptimization(py::module m) {
             py::arg("lb"), py::arg("ub"), cls_doc.AddAccelerationBounds.doc)
         .def("AddJerkBounds", &Class::AddJerkBounds, py::arg("lb"),
             py::arg("ub"), cls_doc.AddJerkBounds.doc)
+        .def("AddEffortBoundsAtNormalizedTimes",
+            &Class::AddEffortBoundsAtNormalizedTimes, py::arg("plant"),
+            py::arg("s"), py::arg("lb") = std::nullopt,
+            py::arg("ub") = std::nullopt, py::arg("plant_context") = nullptr,
+            cls_doc.AddEffortBoundsAtNormalizedTimes.doc)
         .def("AddDurationCost", &Class::AddDurationCost,
             py::arg("weight") = 1.0, cls_doc.AddDurationCost.doc)
         .def("AddPathLengthCost", &Class::AddPathLengthCost,

bindings/pydrake/planning/test/trajectory_optimization_test.py

@@ -32,6 +32,7 @@
     ExpressionConstraint
 )
 from pydrake.multibody.plant import MultibodyPlant
+from pydrake.multibody.parsing import Parser
 import pydrake.solvers as mp
 from pydrake.symbolic import Variable
 from pydrake.systems.framework import InputPortSelection
@@ -254,6 +255,16 @@ def test_kinematic_trajectory_optimization(self):
         trajopt.AddAccelerationBounds(lb=b, ub=b)
         trajopt.AddJerkBounds(lb=b, ub=b)
 
+        plant = MultibodyPlant(time_step=0.0)
+        Parser(plant).AddModelsFromUrl(
+            url="package://drake/examples/multibody/cart_pole/cart_pole.sdf")
+        plant.Finalize()
+        plant_context = plant.CreateDefaultContext()
+        trajopt.AddEffortBoundsAtNormalizedTimes(
+            plant=plant, s=[0.1, 0.9],
+            lb=[-1, -2], ub=[3, 4],
+            plant_context=plant_context)
+
         trajopt.AddDurationCost(weight=1)
         trajopt.AddPathLengthCost(weight=1)
         trajopt.AddPathEnergyCost(weight=1)

planning/trajectory_optimization/BUILD.bazel

@@ -97,6 +97,7 @@ drake_cc_library(
         "//math:bspline_basis",
         "//math:gradient",
         "//math:matrix_util",
+        "//multibody/plant",
         "//solvers:mathematical_program",
         "//solvers:mathematical_program_result",
     ],
@@ -194,6 +195,7 @@ drake_cc_googletest(
     deps = [
         ":kinematic_trajectory_optimization",
         "//common/test_utilities:eigen_matrix_compare",
+        "//multibody/benchmarks/pendulum",
         "//solvers:constraint",
         "//solvers:ipopt_solver",
         "//solvers:osqp_solver",

planning/trajectory_optimization/kinematic_trajectory_optimization.cc

@@ -28,6 +28,8 @@ using math::EigenToStdVector;
 using math::ExtractValue;
 using math::InitializeAutoDiff;
 using math::StdVectorToEigen;
+using multibody::MultibodyForces;
+using multibody::MultibodyPlant;
 using solvers::Binding;
 using solvers::BoundingBoxConstraint;
 using solvers::Constraint;
@@ -187,12 +189,82 @@ class DerivativeConstraint : public Constraint {
   const int derivative_order_;
 };
 
+// Constraint of the form tau_lb <= InverseDynamics(q, v, a) <= tau_ub. The
+// constraint should be bound to the duration followed by all of the control
+// points in the bspline. TODO(russt): The constraint actually only depends on
+// a subset of the control points, depending on s, so we could consider
+// optimizing this aspect.
+class EffortConstraint : public Constraint {
+ public:
+  EffortConstraint(const std::shared_ptr<MultibodyPlant<AutoDiffXd>>& plant,
+                   std::shared_ptr<systems::Context<AutoDiffXd>> plant_context,
+                   const Eigen::VectorXd& lb, const Eigen::VectorXd& ub,
+                   const BsplineTrajectory<double>& bspline, double s)
+      : Constraint(lb.size(), 1 + bspline.num_control_points() * bspline.rows(),
+                   lb, ub),
+        plant_(plant),
+        plant_context_(std::move(plant_context)),
+        num_control_points_(bspline.num_control_points()),
+        s_(s) {
+    // Note: consistency checks on the input arguments already happened in
+    // AddEffortBoundsAtNormalizedTimes().
+    M_q_ = bspline.EvaluateLinearInControlPoints(s, 0).cast<AutoDiffXd>();
+    M_v_ = bspline.EvaluateLinearInControlPoints(s, 1).cast<AutoDiffXd>();
+    M_vdot_ = bspline.EvaluateLinearInControlPoints(s, 2).cast<AutoDiffXd>();
+  }
+
+  void DoEval(const Eigen::Ref<const Eigen::VectorXd>& x,
+              Eigen::VectorXd* y) const override {
+    AutoDiffVecXd y_t;
+    Eval(InitializeAutoDiff(x), &y_t);
+    *y = ExtractValue(y_t);
+  }
+
+  void DoEval(const Eigen::Ref<const AutoDiffVecXd>& x,
+              AutoDiffVecXd* y) const override {
+    y->resize(plant_->num_positions());
+    AutoDiffXd duration = x[0];
+    MatrixX<AutoDiffXd> control_points =
+        x.tail(plant_->num_positions() * num_control_points_)
+            .reshaped(plant_->num_positions(), num_control_points_);
+    plant_context_->SetTime(s_ * duration);
+    plant_->SetPositions(plant_context_.get(), control_points * M_q_);
+    plant_->SetVelocities(plant_context_.get(),
+                          control_points * M_v_ / duration);
+    VectorX<AutoDiffXd> vdot = control_points * M_vdot_ / pow(duration, 2);
+    drake::log()->info("q: {}", fmt_eigen(control_points * M_q_));
+    drake::log()->info("v: {}", fmt_eigen(control_points * M_v_ / duration));
+    drake::log()->info("vdot: {}", fmt_eigen(vdot.transpose()));
+    MultibodyForces<AutoDiffXd> forces_(*plant_);
+    plant_->CalcForceElementsContribution(*plant_context_, &forces_);
+    *y = plant_->CalcInverseDynamics(*plant_context_, vdot, forces_);
+  }
+
+  void DoEval(const Eigen::Ref<const VectorX<symbolic::Variable>>&,
+              VectorX<symbolic::Expression>*) const override {
+    throw std::runtime_error(
+        "EffortConstraint does not support evaluation with Expression.");
+  }
+
+ private:
+  const std::shared_ptr<MultibodyPlant<AutoDiffXd>> plant_;
+  std::shared_ptr<systems::Context<AutoDiffXd>> plant_context_;
+  // M_q_, M_v_, M_vdot_ are matrices such that by multiplying control_points
+  // with these matrices, we get the position q, dqds, and d²qds² respectively.
+  VectorX<AutoDiffXd> M_q_;
+  VectorX<AutoDiffXd> M_v_;
+  VectorX<AutoDiffXd> M_vdot_;
+  int num_control_points_;
+  double s_;
+};
+
 // solvers::LinearConstraint stores linear constraints in the form A*x <= b.
-// When constraints are naturally expressed in the form X*a <= b (the decision
-// variables on the left), it is still a vector constraint, but we must do some
-// work to rewrite it as A*x, and the result is a big sparse A. This method
-// does precisely that. For MatrixXDecisionVariable X and VectorXd a, returns
-// SparseMatrix A and VectorXDecisionVariable x such that Ax == Xa.
+// When constraints are naturally expressed in the form X*a <= b (the
+// decision variables on the left), it is still a vector constraint, but we
+// must do some work to rewrite it as A*x, and the result is a big sparse A.
+// This method does precisely that. For MatrixXDecisionVariable X and
+// VectorXd a, returns SparseMatrix A and VectorXDecisionVariable x such that
+// Ax == Xa.
 std::pair<SparseMatrix<double>, VectorXDecisionVariable> RewriteXa(
     const MatrixXDecisionVariable& X, const VectorXd& a) {
   const int num_nonzeros = (a.array() != 0).count();
@@ -497,6 +569,48 @@ std::vector<Binding<Constraint>> KinematicTrajectoryOptimization::AddJerkBounds(
   return bindings;
 }
 
+std::vector<Binding<Constraint>>
+KinematicTrajectoryOptimization::AddEffortBoundsAtNormalizedTimes(
+    const MultibodyPlant<double>& plant,
+    const Eigen::Ref<const Eigen::VectorXd>& s,
+    const std::optional<Eigen::Ref<const Eigen::VectorXd>>& lb,
+    const std::optional<Eigen::Ref<const Eigen::VectorXd>>& ub,
+    const systems::Context<double>* plant_context) {
+  DRAKE_THROW_UNLESS(s.array().minCoeff() >= 0 && s.array().maxCoeff() <= 1);
+  const MatrixXd B = plant.MakeActuationMatrix();
+  const VectorXd B_tau_lb = B * (lb.value_or(plant.GetEffortLowerLimits()));
+  const VectorXd B_tau_ub = B * (ub.value_or(plant.GetEffortUpperLimits()));
+  DRAKE_THROW_UNLESS(B_tau_lb.size() == num_positions());
+  DRAKE_THROW_UNLESS(B_tau_ub.size() == num_positions());
+  DRAKE_THROW_UNLESS((B_tau_lb.array() <= B_tau_ub.array()).all());
+  if (plant_context != nullptr) {
+    plant.ValidateContext(plant_context);
+  }
+
+  VectorXDecisionVariable vars(1 + num_control_points_ * num_positions_);
+  vars << duration_, control_points_.reshaped(control_points_.size(), 1);
+  std::shared_ptr<MultibodyPlant<AutoDiffXd>> plant_ad(
+      systems::System<double>::ToAutoDiffXd(plant));
+  std::vector<Binding<Constraint>> bindings;
+  for (int i = 0; i < ssize(s); ++i) {
+    // Note: we choose to make a separate context for each s so that the
+    // constraints are thread-safe.
+    std::shared_ptr<systems::Context<AutoDiffXd>> plant_context_ad(
+        plant_ad->CreateDefaultContext());
+    if (plant_context != nullptr) {
+      plant_context_ad->SetTimeStateAndParametersFrom(*plant_context);
+    }
+    bindings.emplace_back(prog_.AddConstraint(
+        std::make_shared<EffortConstraint>(plant_ad,
+                                           std::move(plant_context_ad),
+                                           B_tau_lb, B_tau_ub, bspline_, s[i]),
+        vars));
+    bindings.back().evaluator()->set_description(
+        fmt::format("effort bound at s={}", s[i]));
+  }
+  return bindings;
+}
+
 Binding<LinearCost> KinematicTrajectoryOptimization::AddDurationCost(
     double weight) {
   auto binding = prog_.AddLinearCost(weight * duration_);

planning/trajectory_optimization/kinematic_trajectory_optimization.h

@@ -9,6 +9,7 @@
 
 #include "drake/common/copyable_unique_ptr.h"
 #include "drake/common/trajectories/bspline_trajectory.h"
+#include "drake/multibody/plant/multibody_plant.h"
 #include "drake/solvers/binding.h"
 #include "drake/solvers/mathematical_program.h"
 #include "drake/solvers/mathematical_program_result.h"
@@ -204,6 +205,38 @@ class KinematicTrajectoryOptimization {
       const Eigen::Ref<const Eigen::VectorXd>& lb,
       const Eigen::Ref<const Eigen::VectorXd>& ub);
 
+  /** Adds generic (nonlinear) constraints to enforce the effort limits defined
+  in the plant at a sequence of normalized times, `s`:
+  @verbatim
+  B lb ≤ M(q)v̇ + C(q, v)v - τ_g(q) - τ_app ≤ B ub
+  @endverbatim
+  where q, v, and v̇ are evaluated at s. B is the plant's actuation matrix, and
+  M, C, τ_g, and τ_app are the plant's mass matrix, Coriolis force, gravity,
+  and applied force, respectively. `ub` and `lb` are the upper and lower effort
+  bounds, respectively; if they are not provided then
+  plant.GetEffortLowerLimits() and plant.GetEffortUpperLimits() are used.
+
+  Pass `plant_context` if you have non-default parameters in the context. Note
+  that there are no lifetime requirements on `plant` nsor `plant_context`.
+
+  Note that the convex hull property of the B-splines is not guaranteed to hold
+  here -- effort limits maybe be violated away from the normalized times `s`.
+
+  @pre plant.is_finalized()
+  @pre plant.num_positions() == num_positions()
+  @pre s[i] ∈ [0, 1] for all i
+  @pre B lb ≤ B ub
+
+  @returns A vector of bindings with one effort limit constraint for each `s`.
+  */
+  std::vector<solvers::Binding<solvers::Constraint>>
+  AddEffortBoundsAtNormalizedTimes(
+      const multibody::MultibodyPlant<double>& plant,
+      const Eigen::Ref<const Eigen::VectorXd>& s,
+      const std::optional<Eigen::Ref<const Eigen::VectorXd>>& lb = std::nullopt,
+      const std::optional<Eigen::Ref<const Eigen::VectorXd>>& ub = std::nullopt,
+      const systems::Context<double>* plant_context = nullptr);
+
   /** Adds a linear cost on the duration of the trajectory. */
   solvers::Binding<solvers::LinearCost> AddDurationCost(double weight = 1.0);
 
@@ -234,6 +267,9 @@ class KinematicTrajectoryOptimization {
 
   /* TODO(russt):
   - Support additional (non-convex) costs/constraints on q(t) directly.
+  - AddMultibodyPlantConstraints which adds joint, velocity, acceleration, and
+    effort limits + multibody constraints. Presumably we can't have quaternions
+    in a fully-actuated system.
   */
 
  private:

planning/trajectory_optimization/test/kinematic_trajectory_optimization_test.cc

@@ -10,11 +10,13 @@
 #include "drake/common/test_utilities/eigen_matrix_compare.h"
 #include "drake/common/text_logging.h"
 #include "drake/math/matrix_util.h"
+#include "drake/multibody/benchmarks/pendulum/make_pendulum_plant.h"
 #include "drake/solvers/ipopt_solver.h"
 #include "drake/solvers/osqp_solver.h"
 #include "drake/solvers/solve.h"
 
 using Eigen::MatrixXd;
+using Eigen::Vector2d;
 using Eigen::Vector3d;
 using Eigen::VectorXd;
 using std::nullopt;
@@ -359,6 +361,106 @@ TEST_F(KinematicTrajectoryOptimizationTest, AddJerkBounds) {
   }
 }
 
+GTEST_TEST(KinematicTrajectoryOptimizationMultibodyTest,
+           AddEffortBoundsAtNormalizedTimes) {
+  multibody::benchmarks::pendulum::PendulumParameters parameters;
+  auto plant = multibody::benchmarks::pendulum::MakePendulumPlant(parameters);
+  const int num_control_points = 10;
+  KinematicTrajectoryOptimization trajopt(plant->num_positions(),
+                                          num_control_points);
+
+  const Vector1d tau_lb(-1.0), tau_ub(1.0);
+  auto bindings = trajopt.AddEffortBoundsAtNormalizedTimes(*plant, Vector1d(0),
+                                                           tau_lb, tau_ub);
+  EXPECT_EQ(bindings.size(), 1);
+  EXPECT_EQ(bindings[0].evaluator()->num_constraints(), 1);
+  EXPECT_EQ(bindings[0].evaluator()->lower_bound(), tau_lb);
+  EXPECT_EQ(bindings[0].evaluator()->upper_bound(), tau_ub);
+
+  // Use a small traj opt to compute the (normalized) control points.
+  auto MakeNormalizedTrajectory = [&](double duration, double q0, double v0,
+                                      double vdot0) {
+    KinematicTrajectoryOptimization kto(plant->num_positions(),
+                                        num_control_points);
+    // The spline should be normalized (duration = 1).
+    kto.AddDurationConstraint(1, 1);
+    // We have to undo the time scaling.
+    v0 *= duration;
+    vdot0 *= duration * duration;
+    kto.AddPathPositionConstraint(Vector1d(q0), Vector1d(q0), /* s= */ 0);
+    kto.AddPathVelocityConstraint(Vector1d(v0), Vector1d(v0), /* s= */ 0);
+    kto.AddPathAccelerationConstraint(Vector1d(vdot0), Vector1d(vdot0),
+                                      /* s= */ 0);
+    auto result = Solve(kto.prog());
+    return kto.ReconstructTrajectory(result);
+  };
+
+  // Variables for the constraint are [duration, control_points]
+  VectorXd x(1 + num_control_points);
+  const double duration = 4.23;
+  auto vdot_ulimit = [&](double q, double v) {
+    // ml^2vdot + b*v + mglsin(q) <= tau_ub
+    return (tau_ub[0] - parameters.damping() * v -
+            parameters.m() * parameters.g() * parameters.l() * std::sin(q)) /
+           parameters.m() / parameters.l() / parameters.l();
+  };
+
+  // Initial acceleration below the limit.
+  double q0 = 0, v0 = 0;
+  auto traj = MakeNormalizedTrajectory(duration, q0, v0,
+                                       /* vdot0 = */ 0.9 * vdot_ulimit(q0, v0));
+  x << duration,
+      math::StdVectorToEigen<double>(traj.control_points()).row(0).transpose();
+  EXPECT_TRUE(bindings[0].evaluator()->CheckSatisfied(x));
+
+  // Initial acceleration above the limit.
+  traj = MakeNormalizedTrajectory(duration, q0, v0,
+                                  /* vdot0 = */ 1.1 * vdot_ulimit(q0, v0));
+  x << duration,
+      math::StdVectorToEigen<double>(traj.control_points()).row(0).transpose();
+  EXPECT_FALSE(bindings[0].evaluator()->CheckSatisfied(x));
+
+  // Confirm that the damping is also considered by setting a (large) non-zero
+  // velocity.
+  v0 = -100;
+  traj = MakeNormalizedTrajectory(duration, q0, v0,
+                                  /* vdot0 = */ 0.9 * vdot_ulimit(q0, v0));
+  x << duration,
+      math::StdVectorToEigen<double>(traj.control_points()).row(0).transpose();
+  EXPECT_TRUE(bindings[0].evaluator()->CheckSatisfied(x));
+  traj = MakeNormalizedTrajectory(duration, q0, v0,
+                                  /* vdot0 = */ 1.1 * vdot_ulimit(q0, v0));
+  x << duration,
+      math::StdVectorToEigen<double>(traj.control_points()).row(0).transpose();
+  EXPECT_FALSE(bindings[0].evaluator()->CheckSatisfied(x));
+
+  // Set the damping to zero in the context.
+  auto context = plant->CreateDefaultContext();
+  plant->get_joint(multibody::JointIndex(0))
+      .SetDampingVector(context.get(), Vector1d::Zero());
+  bindings = trajopt.AddEffortBoundsAtNormalizedTimes(
+      *plant, Vector1d(0.5), tau_lb, tau_ub, context.get());
+  EXPECT_EQ(bindings.size(), 1);
+  EXPECT_EQ(bindings[0].evaluator()->num_constraints(), 1);
+  EXPECT_EQ(bindings[0].evaluator()->lower_bound(), tau_lb);
+  EXPECT_EQ(bindings[0].evaluator()->upper_bound(), tau_ub);
+  // The torque that tripped do to high damping is now satisfied.
+  EXPECT_TRUE(bindings[0].evaluator()->CheckSatisfied(x));
+
+  // Multiple normalized times.
+  bindings = trajopt.AddEffortBoundsAtNormalizedTimes(
+      *plant, Vector2d(0.5, 1.0), tau_lb, tau_ub);
+  EXPECT_EQ(bindings.size(), 2);
+
+  // Default bounds.
+  bindings = trajopt.AddEffortBoundsAtNormalizedTimes(*plant, Vector1d(0.25));
+  EXPECT_EQ(bindings.size(), 1);
+  EXPECT_TRUE(CompareMatrices(bindings[0].evaluator()->lower_bound(),
+                              plant->GetEffortLowerLimits(), 1e-14));
+  EXPECT_TRUE(CompareMatrices(bindings[0].evaluator()->upper_bound(),
+                              plant->GetEffortUpperLimits(), 1e-14));
+}
+
 TEST_F(KinematicTrajectoryOptimizationTest, AddDurationCost) {
   EXPECT_EQ(trajopt_.prog().linear_costs().size(), 0);
   auto binding = trajopt_.AddDurationCost(1.0);