chomp_optimizer.cpp

Go to the documentation of this file.

/*********************************************************************
 * Software License Agreement (BSD License)
 *
 *  Copyright (c) 2009, Willow Garage, Inc.
 *  All rights reserved.
 *
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions
 *  are met:
 *
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above
 *     copyright notice, this list of conditions and the following
 *     disclaimer in the documentation and/or other materials provided
 *     with the distribution.
 *   * Neither the name of Willow Garage nor the names of its
 *     contributors may be used to endorse or promote products derived
 *     from this software without specific prior written permission.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 *  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 *  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 *  FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 *  COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 *  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 *  BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 *  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 *  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 *  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 *  ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 *  POSSIBILITY OF SUCH DAMAGE.
 *********************************************************************/
 
/* Author: Mrinal Kalakrishnan */
 
#include <chomp_motion_planner/chomp_optimizer.h>
#include <chomp_motion_planner/chomp_utils.h>
#include <moveit/planning_scene/planning_scene.h>
#include <moveit/robot_model/robot_model.h>
#include <moveit/robot_state/conversions.h>
#include <moveit/utils/logger.hpp>
 
#include <rclcpp/logger.hpp>
#include <rclcpp/logging.hpp>
#include <eigen3/Eigen/Core>
#include <eigen3/Eigen/LU>
#include <random>
#include <visualization_msgs/msg/marker_array.hpp>
 
namespace chomp
{
namespace
{
rclcpp::Logger getLogger()
{
  return moveit::getLogger("chomp_optimizer");
}
}  // namespace
 
ChompOptimizer::ChompOptimizer(ChompTrajectory* trajectory, const planning_scene::PlanningSceneConstPtr& planning_scene,
                               const std::string& planning_group, const ChompParameters* parameters,
                               const moveit::core::RobotState& start_state)
  : full_trajectory_(trajectory)
  , robot_model_(planning_scene->getRobotModel())
  , planning_group_(planning_group)
  , parameters_(parameters)
  , group_trajectory_(*full_trajectory_, planning_group_, DIFF_RULE_LENGTH)
  , planning_scene_(planning_scene)
  , state_(start_state)
  , start_state_(start_state)
  , initialized_(false)
{
  RCLCPP_INFO(getLogger(), "Active collision detector is: %s", planning_scene->getCollisionDetectorName().c_str());
 
  hy_env_ = dynamic_cast<const collision_detection::CollisionEnvHybrid*>(
      planning_scene->getCollisionEnv(planning_scene->getCollisionDetectorName()).get());
  if (!hy_env_)
  {
    RCLCPP_WARN(getLogger(), "Could not initialize hybrid collision world from planning scene");
    return;
  }
 
  initialize();
}
 
void ChompOptimizer::initialize()
{
  // init some variables:
  num_vars_free_ = group_trajectory_.getNumFreePoints();
  num_vars_all_ = group_trajectory_.getNumPoints();
  num_joints_ = group_trajectory_.getNumJoints();
 
  free_vars_start_ = group_trajectory_.getStartIndex();
  free_vars_end_ = group_trajectory_.getEndIndex();
 
  collision_detection::CollisionRequest req;
  collision_detection::CollisionResult res;
  req.group_name = planning_group_;
 
  const auto wt = std::chrono::system_clock::now();
  hy_env_->getCollisionGradients(req, res, state_, &planning_scene_->getAllowedCollisionMatrix(), gsr_);
  RCLCPP_INFO(getLogger(), "First coll check took %f sec",
              std::chrono::duration<double>(std::chrono::system_clock::now() - wt).count());
  num_collision_points_ = 0;
  for (const collision_detection::GradientInfo& gradient : gsr_->gradients_)
  {
    num_collision_points_ += gradient.gradients.size();
  }
 
  // set up the joint costs:
  joint_costs_.reserve(num_joints_);
 
  double max_cost_scale = 0.0;
 
  joint_model_group_ = planning_scene_->getRobotModel()->getJointModelGroup(planning_group_);
 
  const std::vector<const moveit::core::JointModel*> joint_models = joint_model_group_->getActiveJointModels();
  for (size_t i = 0; i < joint_models.size(); ++i)
  {
    double joint_cost = 1.0;
    // nh.param("joint_costs/" + joint_models[i]->getName(), joint_cost, 1.0);
    std::vector<double> derivative_costs(3);
    derivative_costs[0] = joint_cost * parameters_->smoothness_cost_velocity_;
    derivative_costs[1] = joint_cost * parameters_->smoothness_cost_acceleration_;
    derivative_costs[2] = joint_cost * parameters_->smoothness_cost_jerk_;
    joint_costs_.push_back(ChompCost(group_trajectory_, i, derivative_costs, parameters_->ridge_factor_));
    double cost_scale = joint_costs_[i].getMaxQuadCostInvValue();
    if (max_cost_scale < cost_scale)
      max_cost_scale = cost_scale;
  }
 
  // scale the smoothness costs
  for (int i = 0; i < num_joints_; ++i)
  {
    joint_costs_[i].scale(max_cost_scale);
  }
 
  // allocate memory for matrices:
  smoothness_increments_ = Eigen::MatrixXd::Zero(num_vars_free_, num_joints_);
  collision_increments_ = Eigen::MatrixXd::Zero(num_vars_free_, num_joints_);
  final_increments_ = Eigen::MatrixXd::Zero(num_vars_free_, num_joints_);
  smoothness_derivative_ = Eigen::VectorXd::Zero(num_vars_all_);
  jacobian_ = Eigen::MatrixXd::Zero(3, num_joints_);
  jacobian_pseudo_inverse_ = Eigen::MatrixXd::Zero(num_joints_, 3);
  jacobian_jacobian_tranpose_ = Eigen::MatrixXd::Zero(3, 3);
  random_state_ = Eigen::VectorXd::Zero(num_joints_);
  joint_state_velocities_ = Eigen::VectorXd::Zero(num_joints_);
 
  group_trajectory_backup_ = group_trajectory_.getTrajectory();
  best_group_trajectory_ = group_trajectory_.getTrajectory();
 
  collision_point_joint_names_.resize(num_vars_all_, std::vector<std::string>(num_collision_points_));
  collision_point_pos_eigen_.resize(num_vars_all_, EigenSTL::vector_Vector3d(num_collision_points_));
  collision_point_vel_eigen_.resize(num_vars_all_, EigenSTL::vector_Vector3d(num_collision_points_));
  collision_point_acc_eigen_.resize(num_vars_all_, EigenSTL::vector_Vector3d(num_collision_points_));
  joint_axes_.resize(num_vars_all_, EigenSTL::vector_Vector3d(num_joints_));
  joint_positions_.resize(num_vars_all_, EigenSTL::vector_Vector3d(num_joints_));
 
  collision_point_potential_.resize(num_vars_all_, std::vector<double>(num_collision_points_));
  collision_point_vel_mag_.resize(num_vars_all_, std::vector<double>(num_collision_points_));
  collision_point_potential_gradient_.resize(num_vars_all_, EigenSTL::vector_Vector3d(num_collision_points_));
 
  collision_free_iteration_ = 0;
  is_collision_free_ = false;
  state_is_in_collision_.resize(num_vars_all_);
  point_is_in_collision_.resize(num_vars_all_, std::vector<int>(num_collision_points_));
 
  last_improvement_iteration_ = -1;
 
 
  // HMC initialization:
  // momentum_ = Eigen::MatrixXd::Zero(num_vars_free_, num_joints_);
  // random_momentum_ = Eigen::MatrixXd::Zero(num_vars_free_, num_joints_);
  // random_joint_momentum_ = Eigen::VectorXd::Zero(num_vars_free_);
  multivariate_gaussian_.clear();
  stochasticity_factor_ = 1.0;
  for (int i = 0; i < num_joints_; ++i)
  {
    multivariate_gaussian_.push_back(
        MultivariateGaussian(Eigen::VectorXd::Zero(num_vars_free_), joint_costs_[i].getQuadraticCostInverse()));
  }
 
  std::map<std::string, std::string> fixed_link_resolution_map;
  for (int i = 0; i < num_joints_; ++i)
  {
    joint_names_.push_back(joint_model_group_->getActiveJointModels()[i]->getName());
    // RCLCPP_INFO(getLogger(),"Got joint %s", joint_names_[i].c_str());
    registerParents(joint_model_group_->getActiveJointModels()[i]);
    fixed_link_resolution_map[joint_names_[i]] = joint_names_[i];
  }
 
  for (const moveit::core::JointModel* jm : joint_model_group_->getFixedJointModels())
  {
    if (!jm->getParentLinkModel())  // root joint doesn't have a parent
      continue;
 
    fixed_link_resolution_map[jm->getName()] = jm->getParentLinkModel()->getParentJointModel()->getName();
  }
 
  // TODO - is this just the joint_roots_?
  for (const moveit::core::LinkModel* link : joint_model_group_->getUpdatedLinkModels())
  {
    if (fixed_link_resolution_map.find(link->getParentJointModel()->getName()) == fixed_link_resolution_map.end())
    {
      const moveit::core::JointModel* parent_model = nullptr;
      bool found_root = false;
 
      while (!found_root)
      {
        if (parent_model == nullptr)
        {
          parent_model = link->getParentJointModel();
        }
        else
        {
          parent_model = parent_model->getParentLinkModel()->getParentJointModel();
          for (const std::string& joint_name : joint_names_)
          {
            if (parent_model->getName() == joint_name)
            {
              found_root = true;
            }
          }
        }
      }
      fixed_link_resolution_map[link->getParentJointModel()->getName()] = parent_model->getName();
    }
  }
 
  int start = free_vars_start_;
  int end = free_vars_end_;
  for (int i = start; i <= end; ++i)
  {
    size_t j = 0;
    for (const collision_detection::GradientInfo& info : gsr_->gradients_)
    {
      for (size_t k = 0; k < info.sphere_locations.size(); ++k)
      {
        if (fixed_link_resolution_map.find(info.joint_name) != fixed_link_resolution_map.end())
        {
          collision_point_joint_names_[i][j] = fixed_link_resolution_map[info.joint_name];
        }
        else
        {
          RCLCPP_ERROR(getLogger(), "Couldn't find joint %s!", info.joint_name.c_str());
        }
        j++;
      }
    }
  }
  initialized_ = true;
}
 
ChompOptimizer::~ChompOptimizer()
{
  destroy();
}
 
void ChompOptimizer::registerParents(const moveit::core::JointModel* model)
{
  const moveit::core::JointModel* parent_model = nullptr;
  bool found_root = false;
 
  if (model == robot_model_->getRootJoint())
    return;
 
  while (!found_root)
  {
    if (parent_model == nullptr)
    {
      if (model->getParentLinkModel() == nullptr)
      {
        RCLCPP_ERROR(getLogger(), "Model %s not root but has nullptr link model parent", model->getName().c_str());
        return;
      }
      else if (model->getParentLinkModel()->getParentJointModel() == nullptr)
      {
        RCLCPP_ERROR(getLogger(), "Model %s not root but has nullptr joint model parent", model->getName().c_str());
        return;
      }
      parent_model = model->getParentLinkModel()->getParentJointModel();
    }
    else
    {
      if (parent_model == robot_model_->getRootJoint())
      {
        found_root = true;
      }
      else
      {
        parent_model = parent_model->getParentLinkModel()->getParentJointModel();
      }
    }
    joint_parent_map_[model->getName()][parent_model->getName()] = true;
  }
}
 
bool ChompOptimizer::optimize()
{
  bool optimization_result = 0;
 
  const auto start_time = std::chrono::system_clock::now();
  // double averageCostVelocity = 0.0;
  // int currentCostIter = 0;
  int cost_window = 10;
  std::vector<double> costs(cost_window, 0.0);
  // double minimaThreshold = 0.05;
  bool should_break_out = false;
 
  // iterate
  for (iteration_ = 0; iteration_ < parameters_->max_iterations_; ++iteration_)
  {
    performForwardKinematics();
    double c_cost = getCollisionCost();
    double s_cost = getSmoothnessCost();
    double cost = c_cost + s_cost;
 
    RCLCPP_DEBUG(getLogger(), "Collision cost %f, smoothness cost: %f", c_cost, s_cost);
 
 
    // if(parameters_->getAddRandomness() && currentCostIter != -1)
    // {
    //   costs[currentCostIter] = cCost;
    //   currentCostIter++;
 
    //   if(currentCostIter >= costWindow)
    //   {
    //     for(int i = 1; i < costWindow; ++i)
    //     {
    //       averageCostVelocity += (costs.at(i) - costs.at(i - 1));
    //     }
 
    //     averageCostVelocity /= static_cast<double>(costWindow);
    //     currentCostIter = -1;
    //   }
    // }
 
    if (iteration_ == 0)
    {
      best_group_trajectory_ = group_trajectory_.getTrajectory();
      best_group_trajectory_cost_ = cost;
      last_improvement_iteration_ = iteration_;
    }
    else if (cost < best_group_trajectory_cost_)
    {
      best_group_trajectory_ = group_trajectory_.getTrajectory();
      best_group_trajectory_cost_ = cost;
      last_improvement_iteration_ = iteration_;
    }
    calculateSmoothnessIncrements();
    calculateCollisionIncrements();
    calculateTotalIncrements();
 
 
    // if(!parameters_->getUseHamiltonianMonteCarlo())
    // {
    //   // non-stochastic version:
    addIncrementsToTrajectory();
    // }
    // else
    // {
    //   // hamiltonian monte carlo updates:
    //   getRandomMomentum();
    //   updateMomentum();
    //   updatePositionFromMomentum();
    //   stochasticity_factor_ *= parameters_->getHmcAnnealingFactor();
    // }
 
    handleJointLimits();
    updateFullTrajectory();
 
    if (iteration_ % 10 == 0)
    {
      RCLCPP_DEBUG(getLogger(), "iteration: %d", iteration_);
      if (isCurrentTrajectoryMeshToMeshCollisionFree())
      {
        num_collision_free_iterations_ = 0;
        RCLCPP_INFO(getLogger(), "Chomp Got mesh to mesh safety at iter %d. Breaking out early.", iteration_);
        is_collision_free_ = true;
        iteration_++;
        should_break_out = true;
      }
      // } else if(safety == CollisionProximitySpace::InCollisionSafe) {
 
 
      // ROS_DEBUG("Trajectory cost: %f (s=%f, c=%f)", getTrajectoryCost(), getSmoothnessCost(), getCollisionCost());
      // CollisionProximitySpace::TrajectorySafety safety = checkCurrentIterValidity();
      // if(safety == CollisionProximitySpace::MeshToMeshSafe)
      // {
      //   num_collision_free_iterations_ = 0;
      //   RCLCPP_INFO(getLogger(),"Chomp Got mesh to mesh safety at iter %d. Breaking out early.", iteration_);
      //   is_collision_free_ = true;
      //   iteration_++;
      //   should_break_out = true;
      // } else if(safety == CollisionProximitySpace::InCollisionSafe) {
      //   num_collision_free_iterations_ = parameters_->getMaxIterationsAfterCollisionFree();
      //   RCLCPP_INFO(getLogger(),"Chomp Got in collision safety at iter %d. Breaking out soon.", iteration_);
      //   is_collision_free_ = true;
      //   iteration_++;
      //   should_break_out = true;
      // }
      // else
      // {
      //   is_collision_free_ = false;
      // }
    }
 
    if (!parameters_->filter_mode_)
    {
      if (c_cost < parameters_->collision_threshold_)
      {
        num_collision_free_iterations_ = parameters_->max_iterations_after_collision_free_;
        is_collision_free_ = true;
        iteration_++;
        should_break_out = true;
      }
      else
      {
        RCLCPP_DEBUG(getLogger(), "cCost %f over threshold %f", c_cost, parameters_->collision_threshold_);
      }
    }
 
    if (std::chrono::duration<double>(std::chrono::system_clock::now() - start_time).count() >
        parameters_->planning_time_limit_)
    {
      RCLCPP_WARN(getLogger(), "Breaking out early due to time limit constraints.");
      break;
    }
 
 
    // if(fabs(averageCostVelocity) < minimaThreshold && currentCostIter == -1 && !is_collision_free_ &&
    // parameters_->getAddRandomness())
    // {
    //   RCLCPP_INFO(getLogger(),"Detected local minima. Attempting to break out!");
    //   int iter = 0;
    //   bool success = false;
    //   while(iter < 20 && !success)
    //   {
    //     performForwardKinematics();
    //     double original_cost = getTrajectoryCost();
    //     group_trajectory_backup_ = group_trajectory_.getTrajectory();
    //     perturbTrajectory();
    //     handleJointLimits();
    //     updateFullTrajectory();
    //     performForwardKinematics();
    //     double new_cost = getTrajectoryCost();
    //     iter ++;
    //     if(new_cost < original_cost)
    //     {
    //       RCLCPP_INFO(getLogger(),"Got out of minimum in %d iters!", iter);
    //       averageCostVelocity = 0.0;
    //       currentCostIter = 0;
    //       success = true;
    //     }
    //     else
    //     {
    //       group_trajectory_.getTrajectory() = group_trajectory_backup_;
    //       updateFullTrajectory();
    //       currentCostIter = 0;
    //       averageCostVelocity = 0.0;
    //       success = false;
    //     }
 
    //   }
 
    //   if(!success)
    //   {
    //     RCLCPP_INFO(getLogger(),"Failed to exit minimum!");
    //   }
    //}
    // else if (currentCostIter == -1)
    //{
    //  currentCostIter = 0;
    //  averageCostVelocity = 0.0;
    //}
 
    if (should_break_out)
    {
      collision_free_iteration_++;
      if (num_collision_free_iterations_ == 0)
      {
        break;
      }
      else if (collision_free_iteration_ > num_collision_free_iterations_)
      {
        // CollisionProximitySpace::TrajectorySafety safety = checkCurrentIterValidity();
        // if(safety != CollisionProximitySpace::MeshToMeshSafe &&
        //    safety != CollisionProximitySpace::InCollisionSafe) {
        //   ROS_WARN("Apparently regressed");
        // }
        break;
      }
    }
  }
 
  if (is_collision_free_)
  {
    optimization_result = true;
    RCLCPP_INFO(getLogger(), "Chomp path is collision free");
  }
  else
  {
    optimization_result = false;
    RCLCPP_ERROR(getLogger(), "Chomp path is not collision free!");
  }
 
  group_trajectory_.getTrajectory() = best_group_trajectory_;
  updateFullTrajectory();
 
  RCLCPP_INFO(getLogger(), "Terminated after %d iterations, using path from iteration %d", iteration_,
              last_improvement_iteration_);
  RCLCPP_INFO(getLogger(), "Optimization core finished in %f sec",
              std::chrono::duration<double>(std::chrono::system_clock::now() - start_time).count());
  RCLCPP_INFO(getLogger(), "Time per iteration %f sec",
              std::chrono::duration<double>(std::chrono::system_clock::now() - start_time).count() / (iteration_ * 1.0));
 
  return optimization_result;
}
 
bool ChompOptimizer::isCurrentTrajectoryMeshToMeshCollisionFree() const
{
  moveit_msgs::msg::RobotTrajectory traj;
  traj.joint_trajectory.joint_names = joint_names_;
 
  for (size_t i = 0; i < group_trajectory_.getNumPoints(); ++i)
  {
    trajectory_msgs::msg::JointTrajectoryPoint point;
    for (size_t j = 0; j < group_trajectory_.getNumJoints(); ++j)
    {
      point.positions.push_back(best_group_trajectory_(i, j));
    }
    traj.joint_trajectory.points.push_back(point);
  }
  moveit_msgs::msg::RobotState start_state_msg;
  moveit::core::robotStateToRobotStateMsg(start_state_, start_state_msg);
  return planning_scene_->isPathValid(start_state_msg, traj, planning_group_);
}
 
void ChompOptimizer::calculateSmoothnessIncrements()
{
  for (int i = 0; i < num_joints_; ++i)
  {
    joint_costs_[i].getDerivative(group_trajectory_.getJointTrajectory(i), smoothness_derivative_);
    smoothness_increments_.col(i) = -smoothness_derivative_.segment(group_trajectory_.getStartIndex(), num_vars_free_);
  }
}
 
void ChompOptimizer::calculateCollisionIncrements()
{
  double potential;
  double vel_mag_sq;
  double vel_mag;
  Eigen::Vector3d potential_gradient;
  Eigen::Vector3d normalized_velocity;
  Eigen::Matrix3d orthogonal_projector;
  Eigen::Vector3d curvature_vector;
  Eigen::Vector3d cartesian_gradient;
 
  collision_increments_.setZero(num_vars_free_, num_joints_);
 
  int start_point = 0;
  int end_point = free_vars_end_;
 
  // In stochastic descent, simply use a random point in the trajectory, rather than all the trajectory points.
  // This is faster and guaranteed to converge, but it may take more iterations in the worst case.
  if (parameters_->use_stochastic_descent_)
  {
    start_point = static_cast<int>(rsl::uniform_real(0., 1.) * (free_vars_end_ - free_vars_start_) + free_vars_start_);
    if (start_point < free_vars_start_)
      start_point = free_vars_start_;
    if (start_point > free_vars_end_)
      start_point = free_vars_end_;
    end_point = start_point;
  }
  else
  {
    start_point = free_vars_start_;
  }
 
  for (int i = start_point; i <= end_point; ++i)
  {
    for (int j = 0; j < num_collision_points_; ++j)
    {
      potential = collision_point_potential_[i][j];
 
      if (potential < 0.0001)
        continue;
 
      potential_gradient = -collision_point_potential_gradient_[i][j];
 
      vel_mag = collision_point_vel_mag_[i][j];
      vel_mag_sq = vel_mag * vel_mag;
 
      // all math from the CHOMP paper:
 
      normalized_velocity = collision_point_vel_eigen_[i][j] / vel_mag;
      orthogonal_projector = Eigen::Matrix3d::Identity() - (normalized_velocity * normalized_velocity.transpose());
      curvature_vector = (orthogonal_projector * collision_point_acc_eigen_[i][j]) / vel_mag_sq;
      cartesian_gradient = vel_mag * (orthogonal_projector * potential_gradient - potential * curvature_vector);
 
      // pass it through the jacobian transpose to get the increments
      getJacobian(i, collision_point_pos_eigen_[i][j], collision_point_joint_names_[i][j], jacobian_);
 
      if (parameters_->use_pseudo_inverse_)
      {
        calculatePseudoInverse();
        collision_increments_.row(i - free_vars_start_).transpose() -= jacobian_pseudo_inverse_ * cartesian_gradient;
      }
      else
      {
        collision_increments_.row(i - free_vars_start_).transpose() -= jacobian_.transpose() * cartesian_gradient;
      }
 
      /*
        if(point_is_in_collision_[i][j])
        {
        break;
        }
      */
    }
  }
  // cout << collision_increments_ << endl;
}
 
void ChompOptimizer::calculatePseudoInverse()
{
  jacobian_jacobian_tranpose_ =
      jacobian_ * jacobian_.transpose() + Eigen::MatrixXd::Identity(3, 3) * parameters_->pseudo_inverse_ridge_factor_;
  jacobian_pseudo_inverse_ = jacobian_.transpose() * jacobian_jacobian_tranpose_.inverse();
}
 
void ChompOptimizer::calculateTotalIncrements()
{
  for (int i = 0; i < num_joints_; ++i)
  {
    final_increments_.col(i) =
        parameters_->learning_rate_ * (joint_costs_[i].getQuadraticCostInverse() *
                                       (parameters_->smoothness_cost_weight_ * smoothness_increments_.col(i) +
                                        parameters_->obstacle_cost_weight_ * collision_increments_.col(i)));
  }
}
 
void ChompOptimizer::addIncrementsToTrajectory()
{
  const std::vector<const moveit::core::JointModel*>& joint_models = joint_model_group_->getActiveJointModels();
  for (size_t i = 0; i < joint_models.size(); ++i)
  {
    double scale = 1.0;
    double max = final_increments_.col(i).maxCoeff();
    double min = final_increments_.col(i).minCoeff();
    double max_scale = parameters_->joint_update_limit_ / fabs(max);
    double min_scale = parameters_->joint_update_limit_ / fabs(min);
    if (max_scale < scale)
      scale = max_scale;
    if (min_scale < scale)
      scale = min_scale;
    group_trajectory_.getFreeTrajectoryBlock().col(i) += scale * final_increments_.col(i);
  }
  // ROS_DEBUG("Scale: %f",scale);
  // group_trajectory_.getFreeTrajectoryBlock() += scale * final_increments_;
}
 
void ChompOptimizer::updateFullTrajectory()
{
  full_trajectory_->updateFromGroupTrajectory(group_trajectory_);
}
 
void ChompOptimizer::debugCost()
{
  double cost = 0.0;
  for (int i = 0; i < num_joints_; ++i)
    cost += joint_costs_[i].getCost(group_trajectory_.getJointTrajectory(i));
  std::cout << "Cost = " << cost << '\n';
}
 
double ChompOptimizer::getTrajectoryCost()
{
  return getSmoothnessCost() + getCollisionCost();
}
 
double ChompOptimizer::getSmoothnessCost()
{
  double smoothness_cost = 0.0;
  // joint costs:
  for (int i = 0; i < num_joints_; ++i)
    smoothness_cost += joint_costs_[i].getCost(group_trajectory_.getJointTrajectory(i));
 
  return parameters_->smoothness_cost_weight_ * smoothness_cost;
}
 
double ChompOptimizer::getCollisionCost()
{
  double collision_cost = 0.0;
 
  double worst_collision_cost = 0.0;
  worst_collision_cost_state_ = -1;
 
  // collision costs:
  for (int i = free_vars_start_; i <= free_vars_end_; ++i)
  {
    double state_collision_cost = 0.0;
    for (int j = 0; j < num_collision_points_; ++j)
    {
      state_collision_cost += collision_point_potential_[i][j] * collision_point_vel_mag_[i][j];
    }
    collision_cost += state_collision_cost;
    if (state_collision_cost > worst_collision_cost)
    {
      worst_collision_cost = state_collision_cost;
      worst_collision_cost_state_ = i;
    }
  }
 
  return parameters_->obstacle_cost_weight_ * collision_cost;
}
 
void ChompOptimizer::computeJointProperties(int trajectory_point)
{
  for (int j = 0; j < num_joints_; ++j)
  {
    const moveit::core::JointModel* joint_model = state_.getJointModel(joint_names_[j]);
    const moveit::core::RevoluteJointModel* revolute_joint =
        dynamic_cast<const moveit::core::RevoluteJointModel*>(joint_model);
    const moveit::core::PrismaticJointModel* prismatic_joint =
        dynamic_cast<const moveit::core::PrismaticJointModel*>(joint_model);
 
    std::string parent_link_name = joint_model->getParentLinkModel()->getName();
    std::string child_link_name = joint_model->getChildLinkModel()->getName();
    Eigen::Isometry3d joint_transform = state_.getGlobalLinkTransform(parent_link_name) *
                                        (robot_model_->getLinkModel(child_link_name)->getJointOriginTransform() *
                                         (state_.getJointTransform(joint_model)));
 
    // joint_transform = inverseWorldTransform * jointTransform;
    Eigen::Vector3d axis;
 
    if (revolute_joint != nullptr)
    {
      axis = revolute_joint->getAxis();
    }
    else if (prismatic_joint != nullptr)
    {
      axis = prismatic_joint->getAxis();
    }
    else
    {
      axis = Eigen::Vector3d::Identity();
    }
 
    axis = joint_transform * axis;
 
    joint_axes_[trajectory_point][j] = axis;
    joint_positions_[trajectory_point][j] = joint_transform.translation();
  }
}
 
template <typename Derived>
void ChompOptimizer::getJacobian(int trajectory_point, Eigen::Vector3d& collision_point_pos, std::string& joint_name,
                                 Eigen::MatrixBase<Derived>& jacobian) const
{
  for (int j = 0; j < num_joints_; ++j)
  {
    if (isParent(joint_name, joint_names_[j]))
    {
      Eigen::Vector3d column = joint_axes_[trajectory_point][j].cross(
          Eigen::Vector3d(collision_point_pos(0), collision_point_pos(1), collision_point_pos(2)) -
          joint_positions_[trajectory_point][j]);
 
      jacobian.col(j)[0] = column.x();
      jacobian.col(j)[1] = column.y();
      jacobian.col(j)[2] = column.z();
    }
    else
    {
      jacobian.col(j)[0] = 0.0;
      jacobian.col(j)[1] = 0.0;
      jacobian.col(j)[2] = 0.0;
    }
  }
}
 
void ChompOptimizer::handleJointLimits()
{
  const std::vector<const moveit::core::JointModel*> joint_models = joint_model_group_->getActiveJointModels();
  for (size_t joint_i = 0; joint_i < joint_models.size(); ++joint_i)
  {
    const moveit::core::JointModel* joint_model = joint_models[joint_i];
 
    if (joint_model->getType() == moveit::core::JointModel::REVOLUTE)
    {
      const moveit::core::RevoluteJointModel* revolute_joint =
          dynamic_cast<const moveit::core::RevoluteJointModel*>(joint_model);
      if (revolute_joint->isContinuous())
      {
        continue;
      }
    }
 
    const moveit::core::JointModel::Bounds& bounds = joint_model->getVariableBounds();
 
    double joint_max = -DBL_MAX;
    double joint_min = DBL_MAX;
 
    for (const moveit::core::VariableBounds& bound : bounds)
    {
      if (bound.min_position_ < joint_min)
      {
        joint_min = bound.min_position_;
      }
 
      if (bound.max_position_ > joint_max)
      {
        joint_max = bound.max_position_;
      }
    }
 
    int count = 0;
 
    bool violation = false;
    do
    {
      double max_abs_violation = 1e-6;
      double max_violation = 0.0;
      int max_violation_index = 0;
      violation = false;
      for (int i = free_vars_start_; i <= free_vars_end_; ++i)
      {
        double amount = 0.0;
        double absolute_amount = 0.0;
        if (group_trajectory_(i, joint_i) > joint_max)
        {
          amount = joint_max - group_trajectory_(i, joint_i);
          absolute_amount = fabs(amount);
        }
        else if (group_trajectory_(i, joint_i) < joint_min)
        {
          amount = joint_min - group_trajectory_(i, joint_i);
          absolute_amount = fabs(amount);
        }
        if (absolute_amount > max_abs_violation)
        {
          max_abs_violation = absolute_amount;
          max_violation = amount;
          max_violation_index = i;
          violation = true;
        }
      }
 
      if (violation)
      {
        int free_var_index = max_violation_index - free_vars_start_;
        double multiplier =
            max_violation / joint_costs_[joint_i].getQuadraticCostInverse()(free_var_index, free_var_index);
        group_trajectory_.getFreeJointTrajectoryBlock(joint_i) +=
            multiplier * joint_costs_[joint_i].getQuadraticCostInverse().col(free_var_index);
      }
      if (++count > 10)
        break;
    } while (violation);
  }
}
 
void ChompOptimizer::performForwardKinematics()
{
  double inv_time = 1.0 / group_trajectory_.getDiscretization();
  double inv_time_sq = inv_time * inv_time;
 
  // calculate the forward kinematics for the fixed states only in the first iteration:
  int start = free_vars_start_;
  int end = free_vars_end_;
  if (iteration_ == 0)
  {
    start = 0;
    end = num_vars_all_ - 1;
  }
 
  is_collision_free_ = true;
 
  auto total_dur = std::chrono::duration<double>::zero();
 
  // for each point in the trajectory
  for (int i = start; i <= end; ++i)
  {
    // Set Robot state from trajectory point...
    collision_detection::CollisionRequest req;
    collision_detection::CollisionResult res;
    req.group_name = planning_group_;
    setRobotStateFromPoint(group_trajectory_, i);
    auto grad = std::chrono::system_clock::now();
 
    hy_env_->getCollisionGradients(req, res, state_, nullptr, gsr_);
    total_dur += (std::chrono::system_clock::now() - grad);
    computeJointProperties(i);
    state_is_in_collision_[i] = false;
 
    // Keep vars in scope
    {
      size_t j = 0;
      for (const collision_detection::GradientInfo& info : gsr_->gradients_)
      {
        for (size_t k = 0; k < info.sphere_locations.size(); ++k)
        {
          collision_point_pos_eigen_[i][j][0] = info.sphere_locations[k].x();
          collision_point_pos_eigen_[i][j][1] = info.sphere_locations[k].y();
          collision_point_pos_eigen_[i][j][2] = info.sphere_locations[k].z();
 
          collision_point_potential_[i][j] =
              getPotential(info.distances[k], info.sphere_radii[k], parameters_->min_clearance_);
          collision_point_potential_gradient_[i][j][0] = info.gradients[k].x();
          collision_point_potential_gradient_[i][j][1] = info.gradients[k].y();
          collision_point_potential_gradient_[i][j][2] = info.gradients[k].z();
 
          point_is_in_collision_[i][j] = (info.distances[k] - info.sphere_radii[k] < info.sphere_radii[k]);
 
          if (point_is_in_collision_[i][j])
          {
            state_is_in_collision_[i] = true;
            is_collision_free_ = false;
          }
          j++;
        }
      }
    }
  }
 
  // now, get the vel and acc for each collision point (using finite differencing)
  for (int i = free_vars_start_; i <= free_vars_end_; ++i)
  {
    for (int j = 0; j < num_collision_points_; ++j)
    {
      collision_point_vel_eigen_[i][j] = Eigen::Vector3d(0, 0, 0);
      collision_point_acc_eigen_[i][j] = Eigen::Vector3d(0, 0, 0);
      for (int k = -DIFF_RULE_LENGTH / 2; k <= DIFF_RULE_LENGTH / 2; ++k)
      {
        collision_point_vel_eigen_[i][j] +=
            (inv_time * DIFF_RULES[0][k + DIFF_RULE_LENGTH / 2]) * collision_point_pos_eigen_[i + k][j];
        collision_point_acc_eigen_[i][j] +=
            (inv_time_sq * DIFF_RULES[1][k + DIFF_RULE_LENGTH / 2]) * collision_point_pos_eigen_[i + k][j];
      }
 
      // get the norm of the velocity:
      collision_point_vel_mag_[i][j] = collision_point_vel_eigen_[i][j].norm();
    }
  }
}
 
void ChompOptimizer::setRobotStateFromPoint(ChompTrajectory& group_trajectory, int i)
{
  const Eigen::MatrixXd::RowXpr& point = group_trajectory.getTrajectoryPoint(i);
 
  std::vector<double> joint_states;
  joint_states.reserve(group_trajectory.getNumJoints());
  for (size_t j = 0; j < group_trajectory.getNumJoints(); ++j)
    joint_states.emplace_back(point(0, j));
 
  state_.setJointGroupPositions(planning_group_, joint_states);
  state_.update();
}
 
void ChompOptimizer::perturbTrajectory()
{
  // int mid_point = (free_vars_start_ + free_vars_end_) / 2;
  if (worst_collision_cost_state_ < 0)
    return;
  int mid_point = worst_collision_cost_state_;
  moveit::core::RobotState random_state = state_;
  const moveit::core::JointModelGroup* planning_group = state_.getJointModelGroup(planning_group_);
  random_state.setToRandomPositions(planning_group);
  std::vector<double> vals;
  random_state.copyJointGroupPositions(planning_group_, vals);
  double* ptr = &vals[0];
  Eigen::Map<Eigen::VectorXd> random_matrix(ptr, vals.size());
  // Eigen::VectorXd random_matrix = vals;
 
  // convert the state into an increment
  random_matrix -= group_trajectory_.getTrajectoryPoint(mid_point).transpose();
 
  // project the increment orthogonal to joint velocities
  group_trajectory_.getJointVelocities(mid_point, joint_state_velocities_);
  joint_state_velocities_.normalize();
  random_matrix = (Eigen::MatrixXd::Identity(num_joints_, num_joints_) -
                   joint_state_velocities_ * joint_state_velocities_.transpose()) *
                  random_matrix;
 
  int mp_free_vars_index = mid_point - free_vars_start_;
  for (int i = 0; i < num_joints_; ++i)
  {
    group_trajectory_.getFreeJointTrajectoryBlock(i) +=
        joint_costs_[i].getQuadraticCostInverse().col(mp_free_vars_index) * random_state_(i);
  }
}
 
}  // namespace chomp