Skip to content

Commit

Permalink
Merge pull request #15248 from mpetersen94/cubic_pp_sparse_solver
Browse files Browse the repository at this point in the history
trajectories: Switch to Sparse solver for constructing CubicWithContinuousSecondDerivatives
  • Loading branch information
SeanCurtis-TRI committed Jun 29, 2021
2 parents 040c1e3 + 23d2c94 commit 8cc82ec
Show file tree
Hide file tree
Showing 2 changed files with 113 additions and 73 deletions.
167 changes: 99 additions & 68 deletions common/trajectories/piecewise_polynomial.cc
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
#include <memory>
#include <utility>

#include <Eigen/SparseCore>
#include <Eigen/SparseLU>
#include <fmt/format.h>

#include "drake/common/drake_assert.h"
Expand Down Expand Up @@ -810,60 +812,49 @@ int PiecewisePolynomial<T>::
const std::vector<T>& breaks,
const std::vector<MatrixX<T>>& samples,
int row, int col,
MatrixX<T>* A,
std::vector<Eigen::Triplet<T>>* triplet_list,
VectorX<T>* b) {
const std::vector<T>& times = breaks;
const std::vector<MatrixX<T>>& Y = samples;
int N = static_cast<int>(times.size());

DRAKE_DEMAND(A != nullptr);
DRAKE_DEMAND(triplet_list != nullptr);
DRAKE_DEMAND(b != nullptr);
DRAKE_DEMAND(A->rows() == 4 * (N - 1));
DRAKE_DEMAND(A->cols() == 4 * (N - 1));
DRAKE_DEMAND(b->rows() == 4 * (N - 1));
DRAKE_DEMAND(b->rows() == 3 * (N - 1));

int row_idx = 0;
MatrixX<T>& Aref = *A;
std::vector<Eigen::Triplet<T>>& triplet_ref = *triplet_list;
VectorX<T>& bref = *b;

for (int i = 0; i < N - 1; ++i) {
const T dt = times[i + 1] - times[i];

// y_i(x_i) = a0i = Y[i]
Aref(row_idx, 4 * i) = 1;
bref(row_idx++) = Y[i](row, col);

// y_i(x_{i+1}) = y_{i+1}(x_{i}) =>
// a0i + a1i*(x_{i+1} - x_i) + a2i(x_{i+1} - x_i)^2 + a3i(x_{i+1} -
// x_i)^3 = a0{i+1}
Aref(row_idx, 4 * i + 0) = 1;
Aref(row_idx, 4 * i + 1) = dt;
Aref(row_idx, 4 * i + 2) = dt * dt;
Aref(row_idx, 4 * i + 3) = dt * dt * dt;
if (i != N - 2) {
Aref(row_idx++, 4 * (i + 1)) = -1;
} else {
bref(row_idx++) = Y[N - 1](row, col);
}
// Y[i] + a1i*(x_{i+1} - x_i) + a2i(x_{i+1} - x_i)^2 + a3i(x_{i+1} -
// x_i)^3 = Y[i+1]
triplet_ref.push_back(Eigen::Triplet<T>(row_idx, 3 * i + 0, dt));
triplet_ref.push_back(Eigen::Triplet<T>(row_idx, 3 * i + 1, dt * dt));
triplet_ref.push_back(Eigen::Triplet<T>(row_idx, 3 * i + 2, dt * dt * dt));
bref(row_idx++) = Y[i + 1](row, col) - Y[i](row, col);

// y_i'(x_{i+1}) = y_{i+1}'(x_{i}) =>
// a1i + 2*a2i(x_{i+1} - x_i) + 3*a3i(x_{i+1} - x_i)^2 = a1{i+1}
if (i != N - 2) {
Aref(row_idx, 4 * i + 1) = 1;
Aref(row_idx, 4 * i + 2) = 2 * dt;
Aref(row_idx, 4 * i + 3) = 3 * dt * dt;
Aref(row_idx++, 4 * (i + 1) + 1) = -1;
triplet_ref.push_back(Eigen::Triplet<T>(row_idx, 3 * i + 0, 1));
triplet_ref.push_back(Eigen::Triplet<T>(row_idx, 3 * i + 1, 2 * dt));
triplet_ref.push_back(Eigen::Triplet<T>(row_idx, 3 * i + 2, 3 * dt * dt));
triplet_ref.push_back(Eigen::Triplet<T>(row_idx++, 3 * (i + 1), -1));
}

if (i != N - 2) {
// y_i''(x_{i+1}) = y_{i+1}''(x_{i}) =>
// 2*a2i + 6*a3i(x_{i+1} - x_i) = 2*a2{i+1}
Aref(row_idx, 4 * i + 2) = 2;
Aref(row_idx, 4 * i + 3) = 6 * dt;
Aref(row_idx++, 4 * (i + 1) + 2) = -2;
triplet_ref.push_back(Eigen::Triplet<T>(row_idx, 3 * i + 1, 2));
triplet_ref.push_back(Eigen::Triplet<T>(row_idx, 3 * i + 2, 6 * dt));
triplet_ref.push_back(Eigen::Triplet<T>(row_idx++, 3 * (i + 1) + 1, -2));
}
}
DRAKE_DEMAND(row_idx == 4 * (N - 1) - 2);
DRAKE_DEMAND(row_idx == 3 * (N - 1) - 2);
return row_idx;
}

Expand Down Expand Up @@ -902,35 +893,50 @@ PiecewisePolynomial<T>::CubicWithContinuousSecondDerivatives(
polynomials[i].resize(rows, cols);
}

MatrixX<T> A(4 * (N - 1), 4 * (N - 1));
VectorX<T> b(4 * (N - 1));
Eigen::SparseMatrix<T> A(3 * (N - 1), 3 * (N - 1));
VectorX<T> b(3 * (N - 1));
VectorX<T> solution;
VectorX<T> coeffs(4);
Eigen::SparseLU<Eigen::SparseMatrix<T>> solver;

A.setZero();
b.setZero();

// Sets up a linear equation to solve for the coefficients.
for (int j = 0; j < rows; ++j) {
for (int k = 0; k < cols; ++k) {
int row_idx =
SetupCubicSplineInteriorCoeffsLinearSystem(times, Y, j, k, &A, &b);
std::vector<Eigen::Triplet<T>> triplet_list;
// 10 coefficients are needed for the constraints that ensure continuity
// between adjacent segments: 3 for position, 4 for velocity, 3 for
// acceleration. The additional coefficients consist of: 3 for final
// position constraint, 1 for initial velocity constraint & 3 for final
// velocity constraint.
triplet_list.reserve(10 * (N - 2) + 7);
int row_idx = SetupCubicSplineInteriorCoeffsLinearSystem(
times, Y, j, k, &triplet_list, &b);

// Endpoints' velocity matches the given ones.
A(row_idx, 1) = 1;
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 0, 1));
b(row_idx++) = Ydot_start(j, k);

A(row_idx, 4 * (N - 2) + 1) = 1;
A(row_idx, 4 * (N - 2) + 2) = 2 * (times[N - 1] - times[N - 2]);
A(row_idx, 4 * (N - 2) + 3) =
3 * (times[N - 1] - times[N - 2]) * (times[N - 1] - times[N - 2]);
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 3 * (N - 2) + 0, 1));
triplet_list.push_back(Eigen::Triplet<T>(
row_idx, 3 * (N - 2) + 1, 2 * (times[N - 1] - times[N - 2])));
triplet_list.push_back(Eigen::Triplet<T>(
row_idx, 3 * (N - 2) + 2,
3 * (times[N - 1] - times[N - 2]) * (times[N - 1] - times[N - 2])));
b(row_idx++) = Ydot_end(j, k);

// TODO(siyuan.feng): Should switch to a sparse solver.
solution = A.colPivHouseholderQr().solve(b);
A.setFromTriplets(triplet_list.begin(), triplet_list.end());
if (j == 0 && k == 0) {
solver.analyzePattern(A);
}
solver.factorize(A);
solution = solver.solve(b);

for (int i = 0; i < N - 1; ++i) {
polynomials[i](j, k) =
Polynomial<T>(solution.template segment<4>(4 * i));
coeffs(0) = Y[i](j, k);
coeffs.tail(3) = solution.template segment<3>(3 * i);
polynomials[i](j, k) = Polynomial<T>(coeffs);
}
}
}
Expand Down Expand Up @@ -964,66 +970,91 @@ PiecewisePolynomial<T>::CubicWithContinuousSecondDerivatives(
polynomials[i].resize(rows, cols);
}

MatrixX<T> A(4 * (N - 1), 4 * (N - 1));
VectorX<T> b(4 * (N - 1));
Eigen::SparseMatrix<T> A(3 * (N - 1), 3 * (N - 1));
VectorX<T> b(3 * (N - 1));
VectorX<T> solution;
VectorX<T> coeffs(4);
Eigen::SparseLU<Eigen::SparseMatrix<T>> solver;

A.setZero();
b.setZero();

// Sets up a linear equation to solve for the coefficients.
for (int j = 0; j < rows; ++j) {
for (int k = 0; k < cols; ++k) {
int row_idx =
SetupCubicSplineInteriorCoeffsLinearSystem(times, Y, j, k, &A, &b);
std::vector<Eigen::Triplet<T>> triplet_list;
// 10 coefficients are needed for the constraints that ensure continuity
// between adjacent segments: 3 for position, 4 for velocity, 3 for
// acceleration. The additional coefficients consist of: 3 for final
// position constraint, 4 for periodic velocity constraint & 3 for
// periodic acceleration constraint. In the case of "not-a-sample" end
// condition, fewer coefficients are needed.
triplet_list.reserve(10 * (N - 2) + 10);
int row_idx = SetupCubicSplineInteriorCoeffsLinearSystem(
times, Y, j, k, &triplet_list, &b);

if (periodic_end_condition) {
// Time during the last segment.
const T end_dt = times[times.size() - 1] - times[times.size() - 2];
// Enforce velocity between end-of-last and beginning-of-first segments
// is continuous.
A(row_idx, 1) = -1; // Linear term of 1st segment.
A(row_idx, 4 * (N - 2) + 1) = 1; // Linear term of last segment.
A(row_idx, 4 * (N - 2) + 2) = 2 * end_dt; // Squared term of last
// segment.
A(row_idx, 4 * (N - 2) + 3) = 3 * end_dt * end_dt; // Cubic term of
// last segment.
// Linear term of 1st segment.
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 0, -1));
// Linear term of last segment.
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 3 * (N - 2) + 0, 1));
// Squared term of last segment.
triplet_list.push_back(
Eigen::Triplet<T>(row_idx, 3 * (N - 2) + 1, 2 * end_dt));
// Cubic term of last segment.
triplet_list.push_back(
Eigen::Triplet<T>(row_idx, 3 * (N - 2) + 2, 3 * end_dt * end_dt));
b(row_idx++) = 0;

// Enforce that acceleration between end-of-last and beginning-of-first
// segments is continuous.
A(row_idx, 2) = -2; // Quadratic term of 1st segment.
A(row_idx, 4 * (N - 2) + 2) = 2; // Quadratic term of last segment.
A(row_idx, 4 * (N - 2) + 3) = 6 * end_dt; // Cubic term of last
// segment.
// Quadratic term of 1st segment.
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 1, -2));
// Quadratic term of last segment.
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 3 * (N - 2) + 1, 2));
// Cubic term of last segment.
triplet_list.push_back(
Eigen::Triplet<T>(row_idx, 3 * (N - 2) + 2, 6 * end_dt));
b(row_idx++) = 0;
} else {
if (N > 3) {
// Ydddot(times[1]) is continuous.
A(row_idx, 3) = 1; // Cubic term of 1st segment.
A(row_idx, 4 + 3) = -1; // Cubic term of 2nd segment.
// Cubic term of 1st segment.
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 2, 1));
// Cubic term of 2nd segment.
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 3 + 2, -1));
b(row_idx++) = 0;

// Ydddot(times[N-2]) is continuous.
A(row_idx, 4 * (N - 3) + 3) = 1;
A(row_idx, 4 * (N - 2) + 3) = -1;
triplet_list.push_back(
Eigen::Triplet<T>(row_idx, 3 * (N - 3) + 2, 1));
triplet_list.push_back(
Eigen::Triplet<T>(row_idx, 3 * (N - 2) + 2, -1));
b(row_idx++) = 0;
} else {
// Set Jerk to zero if only have 3 points, becomes a quadratic.
A(row_idx, 3) = 1;
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 2, 1));
b(row_idx++) = 0;

A(row_idx, 4 + 3) = 1;
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 3 + 2, 1));
b(row_idx++) = 0;
}
}

// TODO(siyuan.feng): Should switch to a sparse solver.
solution = A.colPivHouseholderQr().solve(b);
A.setFromTriplets(triplet_list.begin(), triplet_list.end());
if (j == 0 && k == 0) {
solver.analyzePattern(A);
}
solver.factorize(A);
solution = solver.solve(b);

for (int i = 0; i < N - 1; ++i) {
polynomials[i](j, k) =
Polynomial<T>(solution.template segment<4>(4 * i));
coeffs(0) = Y[i](j, k);
coeffs.tail(3) = solution.template segment<3>(3 * i);
polynomials[i](j, k) = Polynomial<T>(coeffs);
}
}
}
Expand Down
19 changes: 14 additions & 5 deletions common/trajectories/piecewise_polynomial.h
Original file line number Diff line number Diff line change
Expand Up @@ -815,14 +815,23 @@ class PiecewisePolynomial final : public PiecewiseTrajectory<T> {
// N - 2 acceleration constraints for the interior points:
// Pi''(duration_i) = Pi+1''(0), for i in [0, N - 3]
//
// These sum up to 4 * (N - 1) - 2. This function sets up the above
// constraints. There are still 2 constraints missing, which can be resolved
// by various end point conditions (velocity at the end points /
// "not-a-sample" / etc). These will be specified by the callers.
// The first N - 1 position constraints can be used to directly eliminate the
// constant term of each segment as Pi(0) = a0_i. This reduces the number of
// unknowns to 3 * (N-1) and leaves 3 * (N-1) - 2 constraints. This function
// sets up the remaining constraints. There are still 2 constraints missing,
// which can be resolved by various end point conditions (velocity at the end
// points / "not-a-sample" / etc). These will be specified by the callers.
//
// This function performs the constraint setup for the row and column of the
// sample matrix specified by `row` & `col`. The coefficients for the
// left-hand-side of the constraint equations are prepared for loading into a
// sparse matrix by creating a vector of Triplets, `triplet_list` with each
// representing the row, column and value to add to the sparse matrix. The
// right-hand side of the constraint equations is loaded into `b`.
static int SetupCubicSplineInteriorCoeffsLinearSystem(
const std::vector<T>& breaks,
const std::vector<MatrixX<T>>& samples, int row, int col,
MatrixX<T>* A, VectorX<T>* b);
std::vector<Eigen::Triplet<T>>* triplet_list, VectorX<T>* b);

// Throws std::runtime_error if
// `breaks` and `samples` have different length,
Expand Down

0 comments on commit 8cc82ec

Please sign in to comment.