SpiralSteering.hpp

/**
 * @file SpiralSteering.h
 * @author Olzhas Adiyatov ([email protected])
 * @brief this file implements spline trajectory generation based on these two papers
 *
 * B. Nagy and A. Kelly, “Trajectory generation for car-like robots using cubic curvature polynomials,
 * ”Field and Service Robots, 2001, [Online]. Available:
 * http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.637.356&rep=rep1&type=pdf
 *
 * A. Kelly and B. Nagy, “Reactive Nonholonomic Trajectory Generation via Parametric Optimal Control,”
 *  Int. J. Rob. Res., vol. 22, no. 7–8, pp. 583–601, Jul. 2003, doi: 10.1177/02783649030227008.
 *
 * @date 2021-08-06
 *
 * @copyright Copyright (c) 2021
 *
 */

 /**
  * Acknowledgments: Olzhas Zhumabek, Barry Gilhuly, Akmaral Moldagalieva, Beksultan Rakhim, Denis Fadeyev, Shamak Dutta, Assylbek Dakibay
  * 
  */

#pragma once

#include <gtest/gtest.h>

#include <Eigen/Eigen>
#include <cmath>
#include <cstddef>
#include <fstream>
#include <limits>
#include <memory>
#include <iostream>
#include <functional>
#include "Eigen/src/Core/Matrix.h"
#include "ifopt/bounds.h"

#include <ifopt/problem.h>
#include <ifopt/variable_set.h>
#include <ifopt/constraint_set.h>
#include <ifopt/cost_term.h>
#include <ifopt/ipopt_solver.h>

class Spiral
{
public:
  /**
   * @brief Construct a new Spiral Steering Computation object
   *
   */
  Spiral(size_t pLength = 4) : pLength_{ pLength }, simpsons_resolution_(100)
  {
    ;
  }
  /**
   * @brief compute partial derivatives \frac{\partial x}{\partial q}
   *
   * @param q
   * @return Eigen::VectorXd
   */
  inline Eigen::VectorXd compute_dxdq(const Eigen::Ref<const Eigen::VectorXd>& q) const
  {
    Eigen::VectorXd dxdq;
    dxdq.resize(pLength_ + 1);

    for (size_t i = 0; i < pLength_; ++i)
    {
      dxdq(i) = -1.0 / (i + 1) * S(i + 1, q);
    }

    dxdq(pLength_) = cos(theta(q));

    return dxdq;
  }

  inline Eigen::VectorXd compute_dydq(const Eigen::Ref<const Eigen::VectorXd>& q) const
  {
    // TODO integrate this into compute_grad_g
    Eigen::VectorXd dydq;
    dydq.resize(pLength_ + 1);

    for (size_t i = 0; i < pLength_; ++i)
    {
      dydq(i) = 1.0 / (i + 1) * C(i + 1, q);
    }
    dydq(pLength_) = sin(theta(q));

    return dydq;
  }

  inline Eigen::VectorXd compute_dthetadq(const Eigen::Ref<const Eigen::VectorXd>& q) const
  {
    double s_f = q.tail(1).value();
    Eigen::VectorXd dthetadq;
    dthetadq.resize(pLength_ + 1);
    for (size_t i = 0; i < pLength_; ++i)
    {
      dthetadq(i) = 1.0 / (i + 1.0) * pow(s_f, i + 1);
    }
    dthetadq(pLength_) = kappa(q);
    return dthetadq;
  }

  inline Eigen::VectorXd compute_dkappadq(const Eigen::Ref<const Eigen::VectorXd>& q) const
  {
    double s_f = q.tail(1).value();
    Eigen::VectorXd dkappadq;
    dkappadq.resize(pLength_ + 1);
    for (size_t i = 0; i < pLength_; ++i)
    {
      dkappadq(i) = std::pow(s_f, static_cast<double>(i));
    }
    dkappadq(pLength_) = kappa_prime(q);
    return dkappadq;
  }

  /**
   * @brief polynomial curvature
   *
   * @param q = [a, b, c, d, ..., s_f] = [p, s_f]
   * @return double
   */
  inline double kappa(const Eigen::Ref<const Eigen::VectorXd>& q) const
  {
    double s = q.tail(1).value();
    Eigen::VectorXd p = q.head(q.size() - 1);
    double res = 0;
    for (int i = 0; i < p.size(); ++i)
    {
      // s^i*p(i), i=0, 1, ...
      res += std::pow(s, i) * p(i);
    }
    return res;
  }

  /**
   * @brief derivative of curvature
   *
   * @param q = [a, b, c, d, ..., s_f] = [p, s_f]
   * @return double
   */
  inline double kappa_prime(const Eigen::Ref<const Eigen::VectorXd>& q) const
  {
    double s = q.tail(1).value();
    Eigen::VectorXd p = q.head(q.size() - 1);
    double res = 0;
    for (int i = 1; i < p.size(); ++i)
    {
      // s^(i-1)*p(i), i=1, 2, ...
      res += std::pow(s, i - 1) * p(i) * i;
    }
    return res;
  }
  /**
   * @brief heading angle
   *
   * @param q
   * @return double
   */
  inline double theta(const Eigen::Ref<const Eigen::VectorXd> q) const
  {
    double s = q.tail(1).value();
    Eigen::VectorXd p = q.head(q.size() - 1);
    // theta(s) is integral of kappa(s)
    // p - {a, b, c, d, ...}
    double res = 0;
    for (int i = 0; i < p.size(); ++i)
    {
      // s^(i+1)/(i+1)*p(i)
      // p(0) = a, p(1) = b, ...
      res += p(i) * std::pow(s, i + 1.0) / (i + 1.0);
    }
    return res;
  }
  /**
   * @brief eq. 67
   *
   * @param n
   * @param q
   * @return double
   */
  inline double S(double n, const Eigen::Ref<const Eigen::VectorXd>& q) const
  {
    size_t N = simpsons_resolution_;
    double s = q.tail(1).value();
    Eigen::VectorXd p = q.head(q.size() - 1);
    auto func = [&](double s) -> double {
      Eigen::VectorXd q_temp = q;
      q_temp(q.size() - 1) = s;
      return std::pow(s, n) * sin(theta(q_temp));
    };
    return simpsonRule(N, func, s);
  }
  /**
   * @brief eq. 67
   *
   * @param n
   * @param q
   * @return double
   */
  inline double C(double n, const Eigen::Ref<const Eigen::VectorXd>& q) const
  {
    size_t N = simpsons_resolution_;
    double s = q.tail(1).value();
    auto func = [&](double s) -> double {
      Eigen::VectorXd q_temp = q;
      q_temp(q.size() - 1) = s;
      return std::pow(s, n) * cos(theta(q_temp));
    };
    return simpsonRule(N, func, s);
  };

  inline Eigen::MatrixXd compute_grad_g(const Eigen::Ref<const Eigen::VectorXd>& q) const
  {
    Eigen::VectorXd p = q.head(pLength_);
    const double& s_f = q.tail(1).value();

    Eigen::MatrixXd grad_g;
    grad_g.resize(4, pLength_);

    // a is omitted, b, c, d
    for (int i = 1; i < pLength_; ++i)
    {
      grad_g(0, i - 1) = -1.0 / (i + 1.0) * S(i + 1, q);  // eq. 70
      grad_g(1, i - 1) = 1.0 / (i + 1.0) * C(i + 1, q);   // eq. 70, it has a typo, all terms have positive coeff-s.
    }
    grad_g(0, pLength_ - 1) = cos(theta(q));
    grad_g(1, pLength_ - 1) = sin(theta(q));

    for (int i = 1; i < pLength_; ++i)
    {
      grad_g(2, i - 1) = std::pow(s_f, i + 1) / (i + 1.0);  // s_f^(i+1.0)/(i+1.0)
      grad_g(3, i - 1) = std::pow(s_f, i);                  // s_f^(i)
    }
    grad_g(2, pLength_ - 1) = kappa(q);
    grad_g(3, pLength_ - 1) = kappa_prime(q);

    return grad_g;
  }

  inline Eigen::VectorXd compute_g(const Eigen::Ref<const Eigen::VectorXd>& q,
                                   const Eigen::Ref<const Eigen::Vector4d>& x_b) const
  {
    const double& x_f = x_b(0);
    const double& y_f = x_b(1);
    const double& theta_f = x_b(2);
    const double& kappa_f = x_b(3);

    Eigen::Vector4d g;
    g(0) = C(0, q) - x_f;
    g(1) = S(0, q) - y_f;
    g(2) = theta(q) - theta_f;
    g(3) = kappa(q) - kappa_f;
    return g;
  }

  /**
   * @brief
   *
   * @param n a number of interval for evaluations
   * @param func functor for function evaluation
   * @param s final value of the parameter (boundary?)
   * @return double
   */
  inline double simpsonRule(size_t n, const std::function<double(double)>& func, double s) const
  {
    if (n % 2 == 1)
    {
      std::cerr << "n should be even";
      return std::numeric_limits<double>::quiet_NaN();
    }

    double ds = s / static_cast<double>(n);
    double res = 0;

    res += func(0) + func(s);
    for (int i = 2; i < n; i += 2)
    {
      res += 2.0 * func(i * ds) + 4.0 * func((i - 1) * ds);
    }
    res *= ds / 3.0;
    return res;
  }

  inline size_t getPolynomialDegree() const
  {
    return pLength_;
  }

private:
  size_t simpsons_resolution_;
  size_t pLength_;
};

namespace ifopt
{
using Eigen::Matrix4d;
using Eigen::MatrixXd;
using Eigen::Ref;
using Eigen::Vector4d;
using Eigen::VectorXd;

class ExVariables : public VariableSet
{
public:
  // Every variable set has a name, here "var_set1". this allows the constraints
  // and costs to define values and Jacobians specifically w.r.t this variable set.
  ExVariables(Spiral& spiral, double k0, const Ref<const VectorXd>& q_init)
    : ExVariables("var_set1", spiral, k0, q_init)
  {
  }
  ExVariables(const std::string& name, Spiral& spiral, double k0, const Ref<const VectorXd>& q_init)
    : VariableSet(spiral.getPolynomialDegree() + 1, name), pLength_(spiral.getPolynomialDegree()), k0_(k0)
  {
    // the initial values where the NLP starts iterating from
    // q_ = Eigen::VectorXd::Zero(spiral.getPolynomialDegree() + 1);

    q_ = q_init;
  }

  // Here is where you can transform the Eigen::Vector into whatever
  // internal representation of your variables you have (here two doubles, but
  // can also be complex classes such as splines, etc..
  void SetVariables(const VectorXd& x) override
  {
    q_ = x;
  };

  // Here is the reverse transformation from the internal representation to
  // to the Eigen::Vector
  VectorXd GetValues() const override
  {
    return q_;
  };

  // Each variable has an upper and lower bound set here
  VecBound GetBounds() const override
  {
    VecBound bounds(GetRows());

    bounds.at(0) = Bounds(k0_, k0_);
    for (size_t i = 1; i < pLength_ + 1; ++i)
    {
      bounds.at(i) = NoBound;
    }

    return bounds;
  }

private:
  VectorXd q_;
  double k0_;
  size_t pLength_;
};

class ExConstraint : public ConstraintSet
{
public:
  ExConstraint(double k0, const Eigen::Ref<Eigen::Vector4d>& x_b, Spiral& spiral, double kappa_max)
    : ExConstraint("constraint1", k0, x_b, spiral, kappa_max)
  {
  }

  // This constraint set just contains 1 constraint, however generally
  // each set can contain multiple related constraints.
  ExConstraint(const std::string& name, double k0, const Eigen::Ref<Eigen::Vector4d>& x_b, Spiral& spiral,
               double kappa_max)
    : ConstraintSet(1, name), k0_(k0), x_b_(x_b), spiral_(spiral), kappa_max_(kappa_max)
  {
  }

  // The constraint value minus the constant value "1", moved to bounds.
  VectorXd GetValues() const override
  {
    VectorXd g(GetRows());
    VectorXd q = GetVariables()->GetComponent("var_set1")->GetValues();

    return g;
  };

  // The only constraint in this set is an equality constraint to 1.
  // Constant values should always be put into GetBounds(), not GetValues().
  // For inequality constraints (<,>), use Bounds(x, inf) or Bounds(-inf, x).
  VecBound GetBounds() const override
  {
    VecBound b(GetRows());
    return b;
  }

  // This function provides the first derivative of the constraints.
  // In case this is too difficult to write, you can also tell the solvers to
  // approximate the derivatives by finite differences and not overwrite this
  // function, e.g. in ipopt.cc::use_jacobian_approximation_ = true
  // Attention: see the parent class function for important information on sparsity pattern.
  void FillJacobianBlock(std::string var_set, Jacobian& jac_block) const override
  {
    // must fill only that submatrix of the overall Jacobian that relates
    // to this constraint and "var_set1". even if more constraints or variables
    // classes are added, this submatrix will always start at row 0 and column 0,
    // thereby being independent from the overall problem.
    if (var_set == "var_set1")
    {
      VectorXd q = GetVariables()->GetComponent("var_set1")->GetValues();
      auto dkappadq = spiral_.compute_dkappadq(q);
      for (size_t i = 0; i < dkappadq.size(); ++i)
      {
        jac_block.coeffRef(0, i) = dkappadq(i);
      }
    }
  }

private:
  double k0_;
  Vector4d x_b_;
  Spiral spiral_;
  double kappa_max_;
};

/**
 * @brief Cost term for
 *
 */
class ExCost : public CostTerm
{
public:
  ExCost(double k0, const Ref<const VectorXd>& x_b, Spiral& spiral) : ExCost("cost_term1", k0, x_b, spiral)
  {
  }
  ExCost(const std::string& name, double k0, const Ref<const VectorXd>& x_b, Spiral& spiral)
    : CostTerm(name), k0_(k0), x_b_(x_b), spiral_(spiral)
  {
  }

  double GetCost() const override
  {
    VectorXd q = GetVariables()->GetComponent("var_set1")->GetValues();
    VectorXd g = spiral_.compute_g(q, x_b_);

    g(2) *= L;
    g(3) *= (L * L);

    double cost = 0.5 * g.transpose() * g;
    return cost;
  };
  /**
   * @brief
   *
   * @param var_set
   * @param jac
   */
  void FillJacobianBlock(std::string var_set, Jacobian& jac) const override
  {
    if (var_set == "var_set1")
    {
      VectorXd q = GetVariables()->GetComponent("var_set1")->GetValues();

      VectorXd dxdq = spiral_.compute_dxdq(q);
      VectorXd dydq = spiral_.compute_dydq(q);
      VectorXd dthetadq = spiral_.compute_dthetadq(q);
      VectorXd dkappadq = spiral_.compute_dkappadq(q);

      VectorXd g = spiral_.compute_g(q, x_b_);

      // double L = 0.1;
      MatrixXd K =
          // clang-format off
      (Matrix4d() <<
      1,0,0,0,
      0,1,0,0,
      0,0,L,0,
      0,0,0,L*L
      ).finished();
      // clang-format on

      Vector4d delta_x = K * g;

      for (size_t i = 0; i < q.size(); ++i)
      {
        jac.coeffRef(0, i) =
            dxdq(i) * delta_x(0) + dydq(i) * delta_x(1) + dthetadq(i) * delta_x(2) + dkappadq(i) * delta_x(3);
      }
    }
  }

private:
  Spiral spiral_;
  double k0_;
  Eigen::Vector4d x_b_;
  double L{ 1 };
};

}  // namespace ifopt

/**
 * @brief Class for continuous trajectory generation
 *
 */
class SpiralSteering
{
public:
  /**
   * @brief Construct a new Spiral Steering object
   *
   */
  SpiralSteering()
    : pLength_(4)
    , x_tol_(1e-2)
    , y_tol_(1e-2)
    , theta_tol_(1e-4)
    , kappa_tol_(1e-4)
    , simpsons_resolution_(100)
    , type_(CONSTRAINT_SAT)
    , max_iterations_newtons_method_(1e3)
  {
  }
  /**
   * @brief Set the Optimization option
   *
   */
  void setOptimization()  // TODO rename method name
  {
    type_ = OPTIMIZATION;
  }
  /**
   * @brief Set the Constraint Satisfaction option
   *
   */
  void setConstraintSatisfaction()  // TODO rename method name
  {
    type_ = CONSTRAINT_SAT;
  }

  /**
   * @brief Get a trajectory
   *
   * @param x_b where x_b(0) = x_f, x_b(1) = y_f, x_b(2) = theta_f, x_b(3) = kappa_f
   * @return Eigen::VectorXd
   */
  inline Eigen::VectorXd generate_curve(double k0, const Eigen::Ref<const Eigen::Vector4d>& x_b)
  {
    switch (type_)
    {
      case CONSTRAINT_SAT: {
        return constraintSatisfaction(k0, x_b);
        break;
      };
      case OPTIMIZATION: {
        return optimization(k0, x_b);
        break;
      };
      default:
        std::cerr << "Something went terribly wrong, unknown type of solution method" << std::endl;
        break;
    }
    return {};
  }
  /**
   * @brief Set the Minimum Turn Radius
   *
   * @param R
   */
  void setMinimumTurnRadius(double R)
  {
    kappa_max_ = 1.0 / R;
  }

protected:
  /**
   * @brief
   *
   * @param k0
   * @param x_b
   * @return Eigen::VectorXd
   */
  Eigen::VectorXd optimization(double k0, const Eigen::Ref<const Eigen::Vector4d>& x_b)
  {
    double min_turn_radius = 2.0;
    double kappa_max = 1.0 / min_turn_radius;
    Spiral spiral;

    Eigen::VectorXd q;
    q.resize(pLength_ + 1);

    initializeCoefficients(q, k0, x_b);

    ifopt::Problem nlp;
    nlp.AddVariableSet(std::make_shared<ifopt::ExVariables>(spiral, k0, q));
    // nlp.AddConstraintSet(std::make_shared<ifopt::ExConstraint>(k0, x_b, spiral, kappa_max));
    nlp.AddCostSet(std::make_shared<ifopt::ExCost>(k0, x_b, spiral));

    nlp.PrintCurrent();

    ifopt::IpoptSolver ipopt;
    ipopt.SetOption("linear_solver", "mumps");
    ipopt.SetOption("jacobian_approximation", "exact");
    ipopt.SetOption("print_level", 0);

    ipopt.Solve(nlp);

    Eigen::VectorXd q_star = nlp.GetOptVariables()->GetValues().transpose();
    

    return q_star;
  }

  /**
   * @brief
   *
   * @param k0
   * @param x_b
   * @return Eigen::VectorXd
   */
  Eigen::VectorXd constraintSatisfaction(double k0, const Eigen::Ref<const Eigen::Vector4d>& x_b)
  {
    Spiral spiral;
    const double& x_f = x_b(0);
    const double& y_f = x_b(1);
    const double& theta_f = x_b(2);
    const double& kappa_f = x_b(3);

    Eigen::VectorXd q;
    q.resize(pLength_ + 1);  // a, b, c, d, s_f
    q = Eigen::VectorXd::Zero(pLength_ + 1, 1);

    initializeCoefficients(q, k0, x_b);
    // std::cout << "q0 = " << q.transpose() << "\n" << std::endl;

    std::size_t K = 0;
    std::size_t K_MAX = max_iterations_newtons_method_;

    // std::cout << (q - q_prev).norm() << std::endl;
    Eigen::VectorXd dq = Eigen::VectorXd::Ones(pLength_ + 1, 1);
    dq(0) = 0;

    Eigen::VectorXd g;
    Eigen::MatrixXd grad_g;

    double alpha = 1;

    do
    {
      g = spiral.compute_g(q, x_b);
      grad_g = spiral.compute_grad_g(q);

      // std::cout << g.transpose() << std::endl;
      // std::cout << grad_g << "\n================================" << std::endl;

      // dq.tail(pLength) = grad_g.jacobiSvd(Eigen::ComputeFullU | Eigen::ComputeFullV).solve(-g);
      dq.tail(pLength_) = grad_g.colPivHouseholderQr().solve(-g);

      q = q + alpha * dq;
    } while (toleranceTest(g) && ++K < K_MAX);

    if (K == K_MAX)
    {
      std::cout << "Newton's method was not able to coverge to a solution in " << K_MAX << " iterations." << std::endl;
    }
    return q;
  }

  /**
   * @brief Heuristic initialization for Newton's method
   *
   * B. Nagy and A. Kelly, “Trajectory generation for car-like robots using cubic curvature polynomials,”
   * Field and Service Robots, 2001, [Online].
   * Available: http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.637.356&rep=rep1&type=pdf
   *
   * @param q a parameter vector
   * @param x_b boundary condition, (x_f,y_f,theta_f,kappa_f)
   */

  static void initializeCoefficients(Eigen::Ref<Eigen::VectorXd> q, double k0,
                                     const Eigen::Ref<const Eigen::VectorXd>& x_b)
  {
    double kappa_0 = k0;

    const double& x_f = x_b(0);
    double y_f = x_b(1);
    double theta_f = x_b(2);
    double kappa_f = x_b(3);

    double d = sqrt(x_f * x_f + y_f * y_f);

    double delta_theta = std::abs(theta_f);

    double s_f = d * (delta_theta * delta_theta / 5.0 + 1.0) + 2.0 / 5.0 * delta_theta;
    double c = 0;
    double a = 6 * theta_f / (s_f * s_f) - 2 * kappa_0 / s_f + 4 * kappa_f / s_f;

    double b = 3 / (s_f * s_f) * (kappa_0 + kappa_f) + 6 * theta_f / (s_f * s_f * s_f);

    q(0) = kappa_0;  // a
    q(1) = a;        // b
    q(2) = b;        // c
    q(3) = c;        // d
    q(4) = s_f;
  }

  /**
   * @brief
   *
   * @param g
   * @return true
   * @return false
   */
  bool toleranceTest(const Eigen::Ref<Eigen::Vector4d> g)
  {
    Eigen::Vector4d abs_g = g.cwiseAbs();
    return !(abs_g(0) < x_tol_ && abs_g(1) < y_tol_ && abs_g(2) < theta_tol_ && abs_g(3) < kappa_tol_);
  }

private:
  FRIEND_TEST(SpiralSteeringTest, simpsonsRuleTest);
  FRIEND_TEST(SpiralSteeringTest, gradientCompTest);

  size_t pLength_;

  double x_tol_;
  double y_tol_;
  double theta_tol_;
  double kappa_tol_;

  double kappa_max_;

  size_t simpsons_resolution_;

  size_t max_iterations_newtons_method_;

  enum
  {
    CONSTRAINT_SAT,
    OPTIMIZATION
  } type_;
};