Extracting problemMetrics for SLP

farbod-farshidian · farbod-farshidian · commit 6b43b13e9773 · 2022-12-14T14:45:33.000+01:00
diff --git a/ocs2_slp/include/ocs2_slp/SlpSolver.h b/ocs2_slp/include/ocs2_slp/SlpSolver.h
@@ -65,9 +65,7 @@ class SlpSolver : public SolverBase {
 
   void getPrimalSolution(scalar_t finalTime, PrimalSolution* primalSolutionPtr) const override { *primalSolutionPtr = primalSolution_; }
 
-  const ProblemMetrics& getSolutionMetrics() const override {
-    throw std::runtime_error("[SlpSolver] getSolutionMetrics() not available yet.");
-  }
+  const ProblemMetrics& getSolutionMetrics() const override { return problemMetrics_; }
 
   size_t getNumIterations() const override { return totalNumIterations_; }
 
@@ -124,11 +122,11 @@ class SlpSolver : public SolverBase {
 
   /** Creates QP around t, x(t), u(t). Returns performance metrics at the current {t, x(t), u(t)} */
   PerformanceIndex setupQuadraticSubproblem(const std::vector<AnnotatedTime>& time, const vector_t& initState, const vector_array_t& x,
-                                            const vector_array_t& u);
+                                            const vector_array_t& u, std::vector<Metrics>& metrics);
 
   /** Computes only the performance metrics at the current {t, x(t), u(t)} */
   PerformanceIndex computePerformance(const std::vector<AnnotatedTime>& time, const vector_t& initState, const vector_array_t& x,
-                                      const vector_array_t& u);
+                                      const vector_array_t& u, std::vector<Metrics>& metrics);
 
   /** Returns solution of the QP subproblem in delta coordinates: */
   struct OcpSubproblemSolution {
@@ -143,7 +141,8 @@ class SlpSolver : public SolverBase {
 
   /** Decides on the step to take and overrides given trajectories {x(t), u(t)} <- {x(t) + a*dx(t), u(t) + a*du(t)} */
   slp::StepInfo takeStep(const PerformanceIndex& baseline, const std::vector<AnnotatedTime>& timeDiscretization, const vector_t& initState,
-                         const OcpSubproblemSolution& subproblemSolution, vector_array_t& x, vector_array_t& u);
+                         const OcpSubproblemSolution& subproblemSolution, vector_array_t& x, vector_array_t& u,
+                         std::vector<Metrics>& metrics);
 
   /** Determine convergence after a step */
   slp::Convergence checkConvergence(int iteration, const PerformanceIndex& baseline, const slp::StepInfo& stepInfo) const;
@@ -179,6 +178,9 @@ class SlpSolver : public SolverBase {
   // Iteration performance log
   std::vector<PerformanceIndex> performanceIndeces_;
 
+  // The ProblemMetrics associated to primalSolution_
+  ProblemMetrics problemMetrics_;
+
   // Benchmarking
   size_t numProblems_{0};
   size_t totalNumIterations_{0};
diff --git a/ocs2_slp/src/SlpSolver.cpp b/ocs2_slp/src/SlpSolver.cpp
@@ -35,6 +35,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 #include <ocs2_oc/multiple_shooting/Helpers.h>
 #include <ocs2_oc/multiple_shooting/Initialization.h>
+#include <ocs2_oc/multiple_shooting/MetricsComputation.h>
 #include <ocs2_oc/multiple_shooting/PerformanceIndexComputation.h>
 #include <ocs2_oc/multiple_shooting/Transcription.h>
 #include <ocs2_oc/precondition/Ruzi.h>
@@ -177,6 +178,7 @@ void SlpSolver::runImpl(scalar_t initTime, const vector_t& initState, scalar_t f
 
   // Bookkeeping
   performanceIndeces_.clear();
+  std::vector<Metrics> metrics;
 
   int iter = 0;
   slp::Convergence convergence = slp::Convergence::FALSE;
@@ -186,7 +188,7 @@ void SlpSolver::runImpl(scalar_t initTime, const vector_t& initState, scalar_t f
     }
     // Make QP approximation
     linearQuadraticApproximationTimer_.startTimer();
-    const auto baselinePerformance = setupQuadraticSubproblem(timeDiscretization, initState, x, u);
+    const auto baselinePerformance = setupQuadraticSubproblem(timeDiscretization, initState, x, u, metrics);
     linearQuadraticApproximationTimer_.endTimer();
 
     // Solve LP
@@ -197,7 +199,7 @@ void SlpSolver::runImpl(scalar_t initTime, const vector_t& initState, scalar_t f
 
     // Apply step
     linesearchTimer_.startTimer();
-    const auto stepInfo = takeStep(baselinePerformance, timeDiscretization, initState, deltaSolution, x, u);
+    const auto stepInfo = takeStep(baselinePerformance, timeDiscretization, initState, deltaSolution, x, u, metrics);
     performanceIndeces_.push_back(stepInfo.performanceAfterStep);
     linesearchTimer_.endTimer();
 
@@ -211,6 +213,7 @@ void SlpSolver::runImpl(scalar_t initTime, const vector_t& initState, scalar_t f
 
   computeControllerTimer_.startTimer();
   primalSolution_ = toPrimalSolution(timeDiscretization, std::move(x), std::move(u));
+  problemMetrics_ = multiple_shooting::toProblemMetrics(timeDiscretization, std::move(metrics));
   computeControllerTimer_.endTimer();
 
   ++numProblems_;
@@ -290,7 +293,7 @@ PrimalSolution SlpSolver::toPrimalSolution(const std::vector<AnnotatedTime>& tim
 }
 
 PerformanceIndex SlpSolver::setupQuadraticSubproblem(const std::vector<AnnotatedTime>& time, const vector_t& initState,
-                                                     const vector_array_t& x, const vector_array_t& u) {
+                                                     const vector_array_t& x, const vector_array_t& u, std::vector<Metrics>& metrics) {
   // Problem horizon
   const int N = static_cast<int>(time.size()) - 1;
 
@@ -302,6 +305,7 @@ PerformanceIndex SlpSolver::setupQuadraticSubproblem(const std::vector<Annotated
   stateInputIneqConstraints_.resize(N);
   constraintsProjection_.resize(N);
   projectionMultiplierCoefficients_.resize(N);
+  metrics.resize(N + 1);
 
   std::atomic_int timeIndex{0};
   auto parallelTask = [&](int workerId) {
@@ -314,6 +318,7 @@ PerformanceIndex SlpSolver::setupQuadraticSubproblem(const std::vector<Annotated
       if (time[i].event == AnnotatedTime::Event::PreEvent) {
         // Event node
         auto result = multiple_shooting::setupEventNode(ocpDefinition, time[i].time, x[i], x[i + 1]);
+        metrics[i] = multiple_shooting::computeMetrics(result);
         workerPerformance += multiple_shooting::computePerformanceIndex(result);
         cost_[i] = std::move(result.cost);
         dynamics_[i] = std::move(result.dynamics);
@@ -327,6 +332,7 @@ PerformanceIndex SlpSolver::setupQuadraticSubproblem(const std::vector<Annotated
         const scalar_t ti = getIntervalStart(time[i]);
         const scalar_t dt = getIntervalDuration(time[i], time[i + 1]);
         auto result = multiple_shooting::setupIntermediateNode(ocpDefinition, sensitivityDiscretizer_, ti, dt, x[i], x[i + 1], u[i]);
+        metrics[i] = multiple_shooting::computeMetrics(result);
         workerPerformance += multiple_shooting::computePerformanceIndex(result, dt);
         multiple_shooting::projectTranscription(result, settings_.extractProjectionMultiplier);
         cost_[i] = std::move(result.cost);
@@ -344,6 +350,7 @@ PerformanceIndex SlpSolver::setupQuadraticSubproblem(const std::vector<Annotated
     if (i == N) {  // Only one worker will execute this
       const scalar_t tN = getIntervalStart(time[N]);
       auto result = multiple_shooting::setupTerminalNode(ocpDefinition, tN, x[N]);
+      metrics[i] = multiple_shooting::computeMetrics(result);
       workerPerformance += multiple_shooting::computePerformanceIndex(result);
       cost_[i] = std::move(result.cost);
       stateIneqConstraints_[i] = std::move(result.ineqConstraints);
@@ -365,44 +372,46 @@ PerformanceIndex SlpSolver::setupQuadraticSubproblem(const std::vector<Annotated
 }
 
 PerformanceIndex SlpSolver::computePerformance(const std::vector<AnnotatedTime>& time, const vector_t& initState, const vector_array_t& x,
-                                               const vector_array_t& u) {
-  // Problem horizon
+                                               const vector_array_t& u, std::vector<Metrics>& metrics) {
+  // Problem size
   const int N = static_cast<int>(time.size()) - 1;
+  metrics.resize(N + 1);
 
   std::vector<PerformanceIndex> performance(settings_.nThreads, PerformanceIndex());
   std::atomic_int timeIndex{0};
   auto parallelTask = [&](int workerId) {
     // Get worker specific resources
     OptimalControlProblem& ocpDefinition = ocpDefinitions_[workerId];
-    PerformanceIndex workerPerformance;  // Accumulate performance in local variable
 
     int i = timeIndex++;
     while (i < N) {
       if (time[i].event == AnnotatedTime::Event::PreEvent) {
         // Event node
-        workerPerformance += multiple_shooting::computeEventPerformance(ocpDefinition, time[i].time, x[i], x[i + 1]);
+        metrics[i] = multiple_shooting::computeEventMetrics(ocpDefinition, time[i].time, x[i], x[i + 1]);
+        performance[workerId] += toPerformanceIndex(metrics[i]);
       } else {
         // Normal, intermediate node
         const scalar_t ti = getIntervalStart(time[i]);
         const scalar_t dt = getIntervalDuration(time[i], time[i + 1]);
-        workerPerformance += multiple_shooting::computeIntermediatePerformance(ocpDefinition, discretizer_, ti, dt, x[i], x[i + 1], u[i]);
+        metrics[i] = multiple_shooting::computeIntermediateMetrics(ocpDefinition, discretizer_, ti, dt, x[i], x[i + 1], u[i]);
+        performance[workerId] += toPerformanceIndex(metrics[i], dt);
       }
 
       i = timeIndex++;
     }
 
     if (i == N) {  // Only one worker will execute this
       const scalar_t tN = getIntervalStart(time[N]);
-      workerPerformance += multiple_shooting::computeTerminalPerformance(ocpDefinition, tN, x[N]);
+      metrics[N] = multiple_shooting::computeTerminalMetrics(ocpDefinition, tN, x[N]);
+      performance[workerId] += toPerformanceIndex(metrics[N]);
     }
-
-    // Accumulate! Same worker might run multiple tasks
-    performance[workerId] += workerPerformance;
   };
   runParallel(std::move(parallelTask));
 
-  // Account for init state in performance
-  performance.front().dynamicsViolationSSE += (initState - x.front()).squaredNorm();
+  // Account for initial state in performance
+  const vector_t initDynamicsViolation = initState - x.front();
+  metrics.front().dynamicsViolation += initDynamicsViolation;
+  performance.front().dynamicsViolationSSE += initDynamicsViolation.squaredNorm();
 
   // Sum performance of the threads
   PerformanceIndex totalPerformance = std::accumulate(std::next(performance.begin()), performance.end(), performance.front());
@@ -412,7 +421,7 @@ PerformanceIndex SlpSolver::computePerformance(const std::vector<AnnotatedTime>&
 
 slp::StepInfo SlpSolver::takeStep(const PerformanceIndex& baseline, const std::vector<AnnotatedTime>& timeDiscretization,
                                   const vector_t& initState, const OcpSubproblemSolution& subproblemSolution, vector_array_t& x,
-                                  vector_array_t& u) {
+                                  vector_array_t& u, std::vector<Metrics>& metrics) {
   using StepType = FilterLinesearch::StepType;
 
   /*
@@ -438,13 +447,14 @@ slp::StepInfo SlpSolver::takeStep(const PerformanceIndex& baseline, const std::v
   scalar_t alpha = 1.0;
   vector_array_t xNew(x.size());
   vector_array_t uNew(u.size());
+  std::vector<Metrics> metricsNew(metrics.size());
   do {
     // Compute step
     multiple_shooting::incrementTrajectory(u, du, alpha, uNew);
     multiple_shooting::incrementTrajectory(x, dx, alpha, xNew);
 
     // Compute cost and constraints
-    const PerformanceIndex performanceNew = computePerformance(timeDiscretization, initState, xNew, uNew);
+    const PerformanceIndex performanceNew = computePerformance(timeDiscretization, initState, xNew, uNew, metricsNew);
 
     // Step acceptance and record step type
     bool stepAccepted;
@@ -462,6 +472,7 @@ slp::StepInfo SlpSolver::takeStep(const PerformanceIndex& baseline, const std::v
     if (stepAccepted) {  // Return if step accepted
       x = std::move(xNew);
       u = std::move(uNew);
+      metrics = std::move(metricsNew);
 
       // Prepare step info
       slp::StepInfo stepInfo;
diff --git a/ocs2_sqp/ocs2_sqp/src/SqpSolver.cpp b/ocs2_sqp/ocs2_sqp/src/SqpSolver.cpp
@@ -372,7 +372,6 @@ PerformanceIndex SqpSolver::computePerformance(const std::vector<AnnotatedTime>&
                                                const vector_array_t& u, std::vector<Metrics>& metrics) {
   // Problem size
   const int N = static_cast<int>(time.size()) - 1;
-  metrics.clear();
   metrics.resize(N + 1);
 
   std::vector<PerformanceIndex> performance(settings_.nThreads, PerformanceIndex());
