differential_trajectory_generator.hpp

// Copyright (c) TrajoptLib contributors
 
#pragma once
 
#include <stdint.h>
 
#include <cstddef>
#include <expected>
#include <utility>
#include <vector>
 
#include <sleipnir/autodiff/variable.hpp>
#include <sleipnir/optimization/problem.hpp>
#include <sleipnir/optimization/solver/exit_status.hpp>
 
#include "trajopt/path/path.hpp"
#include "trajopt/path/path_builder.hpp"
#include "trajopt/util/symbol_exports.hpp"
 
namespace trajopt {

struct TRAJOPT_DLLEXPORT DifferentialDrivetrain {
  double mass;

  double moi;

  double wheel_radius;

  double wheel_max_angular_velocity;

  double wheel_max_torque;

  double wheel_cof;

  double trackwidth;
};

struct TRAJOPT_DLLEXPORT DifferentialSolution {
  std::vector<double> dt;

  std::vector<double> x;

  std::vector<double> y;

  std::vector<double> heading;

  std::vector<double> vl;

  std::vector<double> vr;

  std::vector<double> angular_velocity;

  std::vector<double> al;

  std::vector<double> ar;

  std::vector<double> angular_acceleration;

  std::vector<double> Fl;

  std::vector<double> Fr;
};

class TRAJOPT_DLLEXPORT DifferentialTrajectorySample {
 public:
  double timestamp = 0.0;

  double x = 0.0;

  double y = 0.0;

  double heading = 0.0;

  double velocity_l = 0.0;

  double velocity_r = 0.0;

  double angular_velocity = 0.0;

  double acceleration_l = 0.0;

  double acceleration_r = 0.0;

  double angular_acceleration = 0.0;

  double force_l = 0.0;

  double force_r = 0.0;
 
  DifferentialTrajectorySample() = default;

  DifferentialTrajectorySample(double timestamp, double x, double y,
                               double heading, double velocity_l,
                               double velocity_r, double angular_velocity,
                               double acceleration_l, double acceleration_r,
                               double angular_acceleration, double force_l,
                               double force_r)
      : timestamp{timestamp},
        x{x},
        y{y},
        heading{heading},
        velocity_l{velocity_l},
        velocity_r{velocity_r},
        angular_velocity{angular_velocity},
        acceleration_l{acceleration_l},
        acceleration_r{acceleration_r},
        angular_acceleration{angular_acceleration},
        force_l{force_l},
        force_r{force_r} {}
};

class TRAJOPT_DLLEXPORT DifferentialTrajectory {
 public:
  std::vector<DifferentialTrajectorySample> samples;
 
  DifferentialTrajectory() = default;

  explicit DifferentialTrajectory(
      std::vector<DifferentialTrajectorySample> samples)
      : samples{std::move(samples)} {}

  explicit DifferentialTrajectory(const DifferentialSolution& solution) {
    double ts = 0.0;
    for (size_t sample = 0; sample < solution.x.size(); ++sample) {
      samples.emplace_back(
          ts, solution.x[sample], solution.y[sample], solution.heading[sample],
          solution.vl[sample], solution.vr[sample],
          solution.angular_velocity[sample], solution.al[sample],
          solution.ar[sample], solution.angular_acceleration[sample],
          solution.Fl[sample], solution.Fr[sample]);
      ts += solution.dt[sample];
    }
  }
};

using DifferentialPath = Path<DifferentialDrivetrain, DifferentialSolution>;

using DifferentialPathBuilder =
    PathBuilder<DifferentialDrivetrain, DifferentialSolution>;

class TRAJOPT_DLLEXPORT DifferentialTrajectoryGenerator {
 public:
  explicit DifferentialTrajectoryGenerator(DifferentialPathBuilder path_builder,
                                           int64_t handle = 0);

  std::expected<DifferentialSolution, slp::ExitStatus> generate(
      bool diagnostics = false);
 
 private:
  DifferentialPath path;

  std::vector<slp::Variable<double>> x;
  std::vector<slp::Variable<double>> y;
  std::vector<slp::Variable<double>> θ;
  std::vector<slp::Variable<double>> vl;
  std::vector<slp::Variable<double>> vr;
  std::vector<slp::Variable<double>> al;
  std::vector<slp::Variable<double>> ar;

  std::vector<slp::Variable<double>> Fl;
  std::vector<slp::Variable<double>> Fr;

  std::vector<slp::Variable<double>> dts;

  std::vector<size_t> Ns;
 
  slp::Problem<double> problem;
 
  void apply_initial_guess(const DifferentialSolution& solution);
 
  DifferentialSolution construct_differential_solution();
};
 
}  // namespace trajopt