frc971/control_loops/drivetrain/trajectory.cc

#include "frc971/control_loops/drivetrain/trajectory.h"

#include <chrono>

#include "Eigen/Dense"
#include "aos/logging/matrix_logging.h"
#include "frc971/control_loops/dlqr.h"
#include "frc971/control_loops/c2d.h"
#include "frc971/control_loops/drivetrain/distance_spline.h"
#include "frc971/control_loops/drivetrain/drivetrain_config.h"
#include "frc971/control_loops/hybrid_state_feedback_loop.h"
#include "frc971/control_loops/state_feedback_loop.h"

namespace frc971 {
namespace control_loops {
namespace drivetrain {

Trajectory::Trajectory(const DistanceSpline *spline,
                       const DrivetrainConfig<double> &config, double vmax,
                       int num_distance)
    : spline_(spline),
      velocity_drivetrain_(
          ::std::unique_ptr<StateFeedbackLoop<2, 2, 2, double,
                                              StateFeedbackHybridPlant<2, 2, 2>,
                                              HybridKalman<2, 2, 2>>>(
              new StateFeedbackLoop<2, 2, 2, double,
                                    StateFeedbackHybridPlant<2, 2, 2>,
                                    HybridKalman<2, 2, 2>>(
                  config.make_hybrid_drivetrain_velocity_loop()))),
      robot_radius_l_(config.robot_radius),
      robot_radius_r_(config.robot_radius),
      longitudal_acceleration_(3.0),
      lateral_acceleration_(2.0),
      Tlr_to_la_((::Eigen::Matrix<double, 2, 2>() << 0.5, 0.5,
                  -1.0 / (robot_radius_l_ + robot_radius_r_),
                  1.0 / (robot_radius_l_ + robot_radius_r_))
                     .finished()),
      Tla_to_lr_(Tlr_to_la_.inverse()),
      plan_(num_distance, vmax) {}

void Trajectory::LateralAccelPass() {
  for (size_t i = 0; i < plan_.size(); ++i) {
    const double distance = Distance(i);
    plan_[i] = ::std::min(plan_[i], LateralVelocityCurvature(distance));
  }
}

// TODO(austin): Deduplicate this potentially with the backward accel function.
// Need to sort out how the max velocity limit is going to work since the
// velocity and acceleration need to match at all points.
// TODO(austin): Accel check the wheels instead of the center of mass.
double Trajectory::ForwardAcceleration(const double x, const double v) {
  ::Eigen::Matrix<double, 2, 1> K3;
  ::Eigen::Matrix<double, 2, 1> K4;
  ::Eigen::Matrix<double, 2, 1> K5;
  K345(x, &K3, &K4, &K5);

  const ::Eigen::Matrix<double, 2, 1> C = K3 * v * v + K4 * v;
  // Now, solve for all a's and find the best one which meets our criteria.
  double maxa = -::std::numeric_limits<double>::infinity();
  for (const double a : {(voltage_limit_ - C(0, 0)) / K5(0, 0),
                         (voltage_limit_ - C(1, 0)) / K5(1, 0),
                         (-voltage_limit_ - C(0, 0)) / K5(0, 0),
                         (-voltage_limit_ - C(1, 0)) / K5(1, 0)}) {
    const ::Eigen::Matrix<double, 2, 1> U = K5 * a + K3 * v * v + K4 * v;
    if ((U.array().abs() < voltage_limit_ + 1e-6).all()) {
      maxa = ::std::max(maxa, a);
    }
  }

  // Then, assume an acceleration oval and stay inside it.
  const double lateral_acceleration = v * v * spline_->DDXY(x).norm();
  const double ellipse_down_shift = longitudal_acceleration_ * 1.0;
  const double ellipse_width_stretch = ::std::sqrt(
      1.0 / (1.0 - ::std::pow(ellipse_down_shift / (longitudal_acceleration_ +
                                                    ellipse_down_shift),
                              2.0)));
  const double squared =
      1.0 - ::std::pow(lateral_acceleration / lateral_acceleration_ /
                           ellipse_width_stretch,
                       2.0);
  // If we would end up with an imaginary number, cap us at 0 acceleration.
  // TODO(austin): Investigate when this happens, why, and fix it.
  if (squared < 0.0) {
    LOG(ERROR, "Imaginary %f, maxa: %f, fa(%f, %f) -> 0.0\n", squared, maxa, x,
        v);
    return 0.0;
  }
  const double longitudal_acceleration =
      ::std::sqrt(::std::abs(squared)) *
          (longitudal_acceleration_ + ellipse_down_shift) -
      ellipse_down_shift;
  return ::std::min(longitudal_acceleration, maxa);
}

void Trajectory::ForwardPass() {
  plan_[0] = 0.0;
  const double delta_distance = Distance(1) - Distance(0);
  for (size_t i = 0; i < plan_.size() - 1; ++i) {
    const double distance = Distance(i);

    // Integrate our acceleration forward one step.
    plan_[i + 1] = ::std::min(
        plan_[i + 1],
        IntegrateAccelForDistance(
            [this](double x, double v) { return ForwardAcceleration(x, v); },
            plan_[i], distance, delta_distance));
  }
}

double Trajectory::BackwardAcceleration(double x, double v) {
  ::Eigen::Matrix<double, 2, 1> K3;
  ::Eigen::Matrix<double, 2, 1> K4;
  ::Eigen::Matrix<double, 2, 1> K5;
  K345(x, &K3, &K4, &K5);

  // Now, solve for all a's and find the best one which meets our criteria.
  const ::Eigen::Matrix<double, 2, 1> C = K3 * v * v + K4 * v;
  double mina = ::std::numeric_limits<double>::infinity();
  for (const double a : {(voltage_limit_ - C(0, 0)) / K5(0, 0),
                         (voltage_limit_ - C(1, 0)) / K5(1, 0),
                         (-voltage_limit_ - C(0, 0)) / K5(0, 0),
                         (-voltage_limit_ - C(1, 0)) / K5(1, 0)}) {
    const ::Eigen::Matrix<double, 2, 1> U = K5 * a + K3 * v * v + K4 * v;
    if ((U.array().abs() < voltage_limit_ + 1e-6).all()) {
      mina = ::std::min(mina, a);
    }
  }

  // Then, assume an acceleration oval and stay inside it.
  const double lateral_acceleration = v * v * spline_->DDXY(x).norm();
  const double ellipse_down_shift = longitudal_acceleration_ * 1.0;
  const double ellipse_width_stretch = ::std::sqrt(
      1.0 / (1.0 - ::std::pow(ellipse_down_shift / (longitudal_acceleration_ +
                                                    ellipse_down_shift),
                              2.0)));
  const double squared =
      1.0 - ::std::pow(lateral_acceleration / lateral_acceleration_ /
                           ellipse_width_stretch,
                       2.0);
  // If we would end up with an imaginary number, cap us at 0 acceleration.
  // TODO(austin): Investigate when this happens, why, and fix it.
  if (squared < 0.0) {
    LOG(ERROR, "Imaginary %f, mina: %f, fa(%f, %f) -> 0.0\n", squared, mina, x,
        v);
    return 0.0;
  }
  const double longitudal_acceleration =
      -::std::sqrt(::std::abs(squared)) *
          (longitudal_acceleration_ + ellipse_down_shift) +
      ellipse_down_shift;
  return ::std::max(longitudal_acceleration, mina);
}

void Trajectory::BackwardPass() {
  const double delta_distance = Distance(0) - Distance(1);
  plan_.back() = 0.0;
  for (size_t i = plan_.size() - 1; i > 0; --i) {
    const double distance = Distance(i);

    // Integrate our deceleration back one step.
    plan_[i - 1] = ::std::min(
        plan_[i - 1],
        IntegrateAccelForDistance(
            [this](double x, double v) { return BackwardAcceleration(x, v); },
            plan_[i], distance, delta_distance));
  }
}

::Eigen::Matrix<double, 3, 1> Trajectory::FFAcceleration(double distance) {
  size_t before_index;
  size_t after_index;
  if (distance < Distance(1)) {
    // Within the first step.
    after_index = 1;
    // Make sure we don't end up off the beginning of the curve.
    if (distance < 0.0) {
      distance = 0.0;
    }
  } else if (distance > Distance(plan_.size() - 2)) {
    // Within the last step.
    after_index = plan_.size() - 1;
    // Make sure we don't end up off the end of the curve.
    if (distance > length()) {
      distance = length();
    }
  } else {
    // Otherwise do the calculation normally.
    after_index = static_cast<size_t>(
        ::std::ceil(distance / length() * (plan_.size() - 1)));
  }
  before_index = after_index - 1;
  const double before_distance = Distance(before_index);
  const double after_distance = Distance(after_index);

  // Now, compute the velocity that we could have if we accelerated from the
  // previous step and decelerated from the next step.  The min will tell us
  // which is in effect.
  const double velocity_forwards = IntegrateAccelForDistance(
      [this](double x, double v) { return ForwardAcceleration(x, v); },
      plan_[before_index], before_distance, distance - before_distance);
  const double velocity_backward = IntegrateAccelForDistance(
      [this](double x, double v) { return BackwardAcceleration(x, v); },
      plan_[after_index], after_distance, distance - after_distance);

  // And then also make sure we aren't curvature limited.
  const double vcurvature = LateralVelocityCurvature(distance);

  double acceleration;
  double velocity;
  if (vcurvature < velocity_forwards && vcurvature < velocity_backward) {
    // If we are curvature limited, we can't accelerate.
    velocity = vcurvature;
    acceleration = 0.0;
  } else if (velocity_forwards < velocity_backward) {
    // Otherwise, pick the acceleration and velocity from the forward pass if it
    // was the predominate factor in this step.
    velocity = velocity_forwards;
    acceleration = ForwardAcceleration(distance, velocity);
  } else {
    // Otherwise, pick the acceleration and velocity from the backward pass if
    // it was the predominate factor in this step.
    velocity = velocity_backward;
    acceleration = BackwardAcceleration(distance, velocity);
  }
  return (::Eigen::Matrix<double, 3, 1>() << distance, velocity, acceleration)
      .finished();
}

::Eigen::Matrix<double, 2, 1> Trajectory::FFVoltage(double distance) {
  const Eigen::Matrix<double, 3, 1> xva = FFAcceleration(distance);
  const double velocity = xva(1);
  const double acceleration = xva(2);
  const double current_ddtheta = spline_->DDTheta(distance);
  const double current_dtheta = spline_->DTheta(distance);
  // We've now got the equation:
  //     K2 * d^x/dt^2 + K1 (dx/dt)^2 = A * K2 * dx/dt + B * U
  const ::Eigen::Matrix<double, 2, 1> my_K2 = K2(current_dtheta);

  const ::Eigen::Matrix<double, 2, 2> B_inverse =
      velocity_drivetrain_->plant().coefficients().B_continuous.inverse();

  // Now, rephrase it as K5 a + K3 v^2 + K4 v = U
  const ::Eigen::Matrix<double, 2, 1> K5 = B_inverse * my_K2;
  const ::Eigen::Matrix<double, 2, 1> K3 = B_inverse * K1(current_ddtheta);
  const ::Eigen::Matrix<double, 2, 1> K4 =
      -B_inverse * velocity_drivetrain_->plant().coefficients().A_continuous *
      my_K2;

  return K5 * acceleration + K3 * velocity * velocity + K4 * velocity;
}

const ::std::vector<double> Trajectory::Distances() const {
  ::std::vector<double> d;
  d.reserve(plan_.size());
  for (size_t i = 0; i < plan_.size(); ++i) {
    d.push_back(Distance(i));
  }
  return d;
}

::Eigen::Matrix<double, 5, 5> Trajectory::ALinearizedContinuous(
    const ::Eigen::Matrix<double, 5, 1> &state) const {

  const double sintheta = ::std::sin(state(2));
  const double costheta = ::std::cos(state(2));
  const ::Eigen::Matrix<double, 2, 1> linear_angular =
      Tlr_to_la_ * state.block<2, 1>(3, 0);

  // When stopped, just roll with a min velocity.
  double linear_velocity = 0.0;
  constexpr double kMinVelocity = 0.1;
  if (::std::abs(linear_angular(0)) < kMinVelocity / 100.0)  {
    linear_velocity = 0.1;
  } else if (::std::abs(linear_angular(0)) > kMinVelocity) {
    linear_velocity = linear_angular(0);
  } else if (linear_angular(0) > 0) {
    linear_velocity = kMinVelocity;
  } else if (linear_angular(0) < 0) {
    linear_velocity = -kMinVelocity;
  }

  ::Eigen::Matrix<double, 5, 5> result = ::Eigen::Matrix<double, 5, 5>::Zero();
  result(0, 2) = -sintheta * linear_velocity;
  result(0, 3) = 0.5 * costheta;
  result(0, 4) = 0.5 * costheta;

  result(1, 2) = costheta * linear_velocity;
  result(1, 3) = 0.5 * sintheta;
  result(1, 4) = 0.5 * sintheta;

  result(2, 3) = Tlr_to_la_(1, 0);
  result(2, 4) = Tlr_to_la_(1, 1);

  result.block<2, 2>(3, 3) =
      velocity_drivetrain_->plant().coefficients().A_continuous;
  return result;
}

::Eigen::Matrix<double, 5, 2> Trajectory::BLinearizedContinuous() const {
  ::Eigen::Matrix<double, 5, 2> result = ::Eigen::Matrix<double, 5, 2>::Zero();
  result.block<2, 2>(3, 0) =
      velocity_drivetrain_->plant().coefficients().B_continuous;
  return result;
}

void Trajectory::AB(const ::Eigen::Matrix<double, 5, 1> &state,
                    ::std::chrono::nanoseconds dt,
                    ::Eigen::Matrix<double, 5, 5> *A,
                    ::Eigen::Matrix<double, 5, 2> *B) const {
  ::Eigen::Matrix<double, 5, 5> A_linearized_continuous =
      ALinearizedContinuous(state);
  ::Eigen::Matrix<double, 5, 2> B_linearized_continuous =
      BLinearizedContinuous();

  // Now, convert it to discrete.
  controls::C2D(A_linearized_continuous, B_linearized_continuous,
                dt, A, B);
}

::Eigen::Matrix<double, 2, 5> Trajectory::KForState(
    const ::Eigen::Matrix<double, 5, 1> &state, ::std::chrono::nanoseconds dt,
    const ::Eigen::DiagonalMatrix<double, 5> &Q,
    const ::Eigen::DiagonalMatrix<double, 2> &R) const {
  ::Eigen::Matrix<double, 5, 5> A;
  ::Eigen::Matrix<double, 5, 2> B;
  AB(state, dt, &A, &B);

  ::Eigen::Matrix<double, 5, 5> S = ::Eigen::Matrix<double, 5, 5>::Zero();
  ::Eigen::Matrix<double, 2, 5> K = ::Eigen::Matrix<double, 2, 5>::Zero();

  int info = ::frc971::controls::dlqr<5, 2>(A, B, Q, R, &K, &S);
  if (info == 0) {
    LOG_MATRIX(INFO, "K", K);
  } else {
    LOG(ERROR, "Failed to solve %d, controllability: %d\n", info,
        controls::Controllability(A, B));
    // TODO(austin): Can we be more clever here?  Use the last one?  We should
    // collect more info about when this breaks down from logs.
    K = ::Eigen::Matrix<double, 2, 5>::Zero();
  }
  ::Eigen::EigenSolver<::Eigen::Matrix<double, 5, 5>> eigensolver(A - B * K);
  const auto eigenvalues = eigensolver.eigenvalues();
  LOG(DEBUG,
      "Eigenvalues: (%f + %fj), (%f + %fj), (%f + %fj), (%f + %fj), (%f + "
      "%fj)\n",
      eigenvalues(0).real(), eigenvalues(0).imag(), eigenvalues(1).real(),
      eigenvalues(1).imag(), eigenvalues(2).real(), eigenvalues(2).imag(),
      eigenvalues(3).real(), eigenvalues(3).imag(), eigenvalues(4).real(),
      eigenvalues(4).imag());
  return K;
}

const ::Eigen::Matrix<double, 5, 1> Trajectory::GoalState(double distance,
                                                          double velocity) {
  ::Eigen::Matrix<double, 5, 1> result;
  result.block<2, 1>(0, 0) = spline_->XY(distance);
  result(2, 0) = spline_->Theta(distance);

  result.block<2, 1>(3, 0) = Tla_to_lr_ *
                             (::Eigen::Matrix<double, 2, 1>() << velocity,
                              spline_->DThetaDt(distance, velocity))
                                 .finished();
  return result;
}

::std::vector<::Eigen::Matrix<double, 3, 1>> Trajectory::PlanXVA(
    ::std::chrono::nanoseconds dt) {
  double dt_float =
      ::std::chrono::duration_cast<::std::chrono::duration<double>>(dt).count();
  double t = 0.0;
  ::Eigen::Matrix<double, 2, 1> state = ::Eigen::Matrix<double, 2, 1>::Zero();

  ::std::vector<::Eigen::Matrix<double, 3, 1>> result;
  result.emplace_back(FFAcceleration(0));
  result.back()(1) = 0.0;

  while (state(0) < length() - 1e-4) {
    t += dt_float;

    // TODO(austin): This feels like something that should be pulled out into
    // a library for re-use.
    state = RungeKutta(
        [this](const ::Eigen::Matrix<double, 2, 1> x) {
          ::Eigen::Matrix<double, 3, 1> xva = FFAcceleration(x(0));
          return (::Eigen::Matrix<double, 2, 1>() << x(1), xva(2)).finished();
        },
        state, dt_float);

    result.emplace_back(FFAcceleration(state(0)));
    state(1) = result.back()(1);
  }
  return result;
}

}  // namespace drivetrain
}  // namespace control_loops
}  // namespace frc971