residuals.h

Go to the documentation of this file.

#include <algorithm>
#include <cmath>
#include <iostream>
#include <istream>
#include <vector>
 
#include <ArmarXCore/core/logging/Logging.h>
 
#include <armarx/navigation/algorithms/orientation_aware/smoothing/Costmap3DWrapper.h>
#include <armarx/navigation/algorithms/orientation_aware/smoothing/io.h>
 
#include <ceres/ceres.h>
 
namespace armarx::navigation::algorithms::orientation_aware::smoothing
{
    // Used ChatGPT for initial draft
    // Full CHOMP-like smoother using Ceres with spacing residual and correct angle handling.
    //
    // NOTE: This is a compact but complete example. Adapt the centerToRobot transform, SDF grid source,
    // weights, and solver options to your codebase.
 
 
    // -----------------------------
    // Data structures
    // -----------------------------
    struct CenterPoint
    {
        double x, y, theta, v;
    };
 
    // -----------------------------
    // Utilities: angle diff + normalize
    // -----------------------------
    template <typename T>
    T
    angleDiff(const T& a, const T& b)
    {
        // returns normalized a - b in [-pi, pi]
        T diff = a - b;
        return ceres::atan2(ceres::sin(diff), ceres::cos(diff));
    }
 
    inline double
    normalizeAngle(double a)
    {
        return std::atan2(std::sin(a), std::cos(a));
    }
 
    static void
    normalizeTrajectoryAngles(std::vector<CenterPoint>& traj)
    {
        for (auto& p : traj)
            p.theta = normalizeAngle(p.theta);
    }
 
    // -----------------------------
    // Residuals
    // -----------------------------
 
    // Pose second-difference smoothness residual (x,y,theta)
    struct PoseSmoothResidual
    {
        PoseSmoothResidual(double pos_weight = 1.0, double orientation_weight = 1.0) :
            pos_w(pos_weight), ori_w(orientation_weight)
        {
        }
 
        template <typename T>
        bool
        operator()(const T* const c_prev,
                   const T* const c_i,
                   const T* const c_next,
                   T* residual) const
        {
            residual[0] = T(pos_w) * (c_prev[0] - T(2.0) * c_i[0] + c_next[0]); // x
            residual[1] = T(pos_w) * (c_prev[1] - T(2.0) * c_i[1] + c_next[1]); // y
            // theta second-diff using angleDiff to handle wrap
            T dprev = angleDiff(c_prev[2], c_i[2]);
            T dnext = angleDiff(c_i[2], c_next[2]);
            residual[2] = T(ori_w) * (dprev - dnext);
            return true;
        }
 
        double pos_w;
        double ori_w;
    };
 
    // Residual to enforce smooth orientation transition at start and end of trajectory
    struct StartRotationResidual
    {
        StartRotationResidual(double w) : w_(w)
        {
        }
 
        template <typename T>
        bool
        operator()(const T* const theta0, const T* const theta1, T* residuals) const
        {
 
            T d = angleDiff(theta1[2], theta0[2]);
            residuals[0] = T(w_) * d;
            return true;
        }
 
        double w_;
    };
 
    // Velocity (v) second-difference smoothness
    struct VelSmoothResidual
    {
        VelSmoothResidual(double weight = 1.0) : w(weight)
        {
        }
 
        template <typename T>
        bool
        operator()(const T* const c_prev,
                   const T* const c_i,
                   const T* const c_next,
                   T* residual) const
        {
            residual[0] = T(w) * (c_prev[3] - T(2.0) * c_i[3] + c_next[3]); // v
            return true;
        }
 
        double w;
    };
 
    // Jerk residual (optional) for tighter smoothness on pose (uses 4 consecutive centers)
    struct PoseJerkResidual
    {
        PoseJerkResidual(double weight = 1.0) : w(weight)
        {
        }
 
        template <typename T>
        bool
        operator()(const T* const c_m2,
                   const T* const c_m1,
                   const T* const c_i,
                   const T* const c_p1,
                   T* residual) const
        {
            residual[0] = T(w) * (c_m2[0] - T(3.0) * c_m1[0] + T(3.0) * c_i[0] - c_p1[0]);
            residual[1] = T(w) * (c_m2[1] - T(3.0) * c_m1[1] + T(3.0) * c_i[1] - c_p1[1]);
            // approximate angle jerk (small-angle assumption)
            T a_m2 = c_m2[2], a_m1 = c_m1[2], a_i = c_i[2], a_p1 = c_p1[2];
            T r = angleDiff(a_m2, a_i) - T(3.0) * angleDiff(a_m1, a_i) +
                  T(3.0) * angleDiff(a_i, a_i) - angleDiff(a_p1, a_i);
            residual[2] = T(w) * r;
            // T d1 = angleDiff(a_m1, a_i);
            // T d2 = angleDiff(a_i,   a_p1);
            // T d3 = angleDiff(a_p1,  a_p2);
            //
            // T jerk = d1 - T(2.0)*d2 + d3;
            // residual[2] = T(w) * jerk;
 
            return true;
        }
 
        double w;
    };
 
    // obstacle residual using Costmap3DWrapper; weighted
    struct ObstacleResidual
    {
        ObstacleResidual(const Costmap3DWrapper& wrapper,
                         double weight,
                         double max_dist,
                         double collision_residual) :
            wrapper_(wrapper),
            w_(weight),
            max_dist_(max_dist),
            collision_residual_(collision_residual)
        {
        }
 
        template <typename T>
        T
        obstacleCost(const T& distance_m) const
        {
            const T k = T(0.03); // decay factor
            const T max_cost = T(collision_residual_);
 
            if (distance_m < T(0))
            {
                // Inside obstacle → constant high cost, no gradient surprise
                return max_cost;
            }
 
            if (distance_m > T(1.0))
            {
                // Beyond 1m → 0 influence
                return T(0);
            }
 
            return max_cost * ceres::exp(-k * distance_m);
        }
 
        template <typename T>
        bool
        operator()(const T* const node, T* residual) const
        {
            T obs_dist;
            wrapper_(&node[0], &node[1], &node[2], &obs_dist);
            // residual[0] = obstacleCost(obs_dist);
            // const T max_dist = T(max_dist_);
            if (obs_dist > max_dist_)
            {
                residual[0] = T(0.0);
            }
            else if (obs_dist < T(0.0))
            {
                // residual[0] = T(collision_residual_);
                const auto tmp = ((max_dist_ - obs_dist) / max_dist_);
                // residual[0] = T(w_) * T(collision_residual_) * tmp;
                residual[0] = T(w_) * T(w_) * (tmp + T(collision_residual_));
            }
            else
            {
                residual[0] = T(w_) * ((max_dist_ - obs_dist) / max_dist_);
            }
            //
            return true;
        }
 
        const Costmap3DWrapper& wrapper_;
        double w_;
        double max_dist_;
        double collision_residual_;
    };
 
    // Tracking residual to original center traj: (x,y,theta,v)
    struct TrackingResidual
    {
        TrackingResidual(double xr,
                         double yr,
                         double thetar,
                         double vr,
                         double wx,
                         double wy,
                         double wth,
                         double wv) :
            xr_(xr), yr_(yr), thetar_(thetar), vr_(vr), wx_(wx), wy_(wy), wth_(wth), wv_(wv)
        {
        }
 
        template <typename T>
        bool
        operator()(const T* const c_i, T* residual) const
        {
            residual[0] = T(wx_) * (c_i[0] - T(xr_));
            residual[1] = T(wy_) * (c_i[1] - T(yr_));
            residual[2] = T(wth_) * angleDiff(c_i[2], T(thetar_));
            residual[3] = T(wv_) * (c_i[3] - T(vr_));
            return true;
        }
 
        double xr_, yr_, thetar_, vr_;
        double wx_, wy_, wth_, wv_;
    };
 
    // Velocity limit hinge residual: penalize v > vmax
    struct VelLimitResidual
    {
        VelLimitResidual(double vmax, double weight = 1.0) : vmax_(vmax), w(weight)
        {
        }
 
        template <typename T>
        bool
        operator()(const T* const c_i, T* residual) const
        {
            T v = c_i[3];
            T diff = v - T(vmax_);
            if (diff > T(0))
                residual[0] = T(w) * diff;
            else
                residual[0] = T(0.0);
            return true;
        }
 
        double vmax_;
        double w;
    };
 
    // Robot smoothness: second-diff applied to robot pose derived from center.
    // The mapping center->robot is templated for AutoDiff.
    struct RobotSmoothResidual
    {
        RobotSmoothResidual(double weight = 1.0, double dx = -0.5, double dy = 0.0) :
            w(weight), dx_(dx), dy_(dy)
        {
        }
 
        template <typename T>
        inline void
        centerToRobotTemplated(const T* const c, T& rx, T& ry, T& rtheta) const
        {
            T cx = c[0], cy = c[1], cth = c[2];
            T ccos = ceres::cos(cth), csin = ceres::sin(cth);
            rx = cx + ccos * T(dx_) - csin * T(dy_);
            ry = cy + csin * T(dx_) + ccos * T(dy_);
            rtheta = cth; // robot orientation equals center orientation here; adapt if needed
        }
 
        template <typename T>
        bool
        operator()(const T* const c_prev,
                   const T* const c_i,
                   const T* const c_next,
                   T* residual) const
        {
            T rpx, rpy, rpth, rix, riy, rith, rnx, rny, rnth;
            centerToRobotTemplated(c_prev, rpx, rpy, rpth);
            centerToRobotTemplated(c_i, rix, riy, rith);
            centerToRobotTemplated(c_next, rnx, rny, rnth);
            residual[0] = T(w) * (rpx - T(2.0) * rix + rnx);
            residual[1] = T(w) * (rpy - T(2.0) * riy + rny);
            // Proper theta second-diff via angleDiff:
            residual[2] = T(w) * (angleDiff(rpth, rith) - angleDiff(rnth, rith));
            return true;
        }
 
        double w;
        double dx_, dy_;
    };
 
    // Spacing residual: keeps |p_{i+1} - p_i| close to average spacing d_avg
    struct SpacingResidual
    {
        SpacingResidual(double d_avg, double weight) : d_avg_(d_avg), w(weight)
        {
        }
 
        template <typename T>
        bool
        operator()(const T* const c_i, const T* const c_ip1, T* residual) const
        {
            T dx = c_ip1[0] - c_i[0];
            T dy = c_ip1[1] - c_i[1];
            T dist = ceres::sqrt(dx * dx + dy * dy);
            residual[0] = T(w) * (dist - T(d_avg_));
            return true;
        }
 
        double d_avg_;
        double w;
    };
 
    // -----------------------------
    // Options & helpers
    // -----------------------------
    static double
    segment_length(const CenterPoint& a, const CenterPoint& b)
    {
        return std::hypot(b.x - a.x, b.y - a.y);
    }
 
    static double
    computeAverageSpacing(const std::vector<CenterPoint>& traj)
    {
        if (traj.size() < 2)
            return 0.0;
        double sum = 0.0;
        for (size_t i = 0; i + 1 < traj.size(); ++i)
            sum += segment_length(traj[i], traj[i + 1]);
        return sum / double(traj.size() - 1);
    }
 
    // -----------------------------
    // Main optimizer
    // -----------------------------
    void
    optimizeTrajectoryCeres(std::vector<CenterPoint>& traj,
                            const Costmap3DWrapper& costmap_wrapper,
                            const std::vector<CenterPoint>& traj_original,
                            const io::SmoothingParams& opts)
    {
        const int N = static_cast<int>(traj.size());
        if (N < 4)
        {
            std::cerr << "[optimizeTrajectoryCeres] need at least 4 points\n";
            return;
        }
 
        // Normalize input angles
        normalizeTrajectoryAngles(traj);
 
        ceres::Problem problem;
        // Add parameter blocks
        for (int i = 0; i < N; ++i)
            problem.AddParameterBlock(&traj[i].x, 4);
 
        // Hard anchors: fix start and goal
        problem.SetParameterBlockConstant(&traj.front().x);
        problem.SetParameterBlockConstant(&traj.back().x);
 
        // 1) Pose smoothness  (i = 1 .. N-2)
        for (int i = 1; i <= N - 2; ++i)
        {
            problem.AddResidualBlock(
                new ceres::AutoDiffCostFunction<PoseSmoothResidual, 3, 4, 4, 4>(
                    new PoseSmoothResidual(opts.w_pose_smooth, opts.w_pose_smooth_ori)),
                nullptr,
                &traj[i - 1].x,
                &traj[i].x,
                &traj[i + 1].x);
        }
 
        // Start and end rotation smoothness
        {
            // Start boundary
            problem.AddResidualBlock(
                new ceres::AutoDiffCostFunction<StartRotationResidual, 1, 4, 4>(
                    new StartRotationResidual(opts.w_boundary)),
                nullptr,
                &traj[0].x,
                &traj[1].x);
 
            // End boundary
            int N = traj.size();
            problem.AddResidualBlock(
                new ceres::AutoDiffCostFunction<StartRotationResidual, 1, 4, 4>(
                    new StartRotationResidual(opts.w_boundary)),
                nullptr,
                &traj[N - 2].x,
                &traj[N - 1].x);
        }
 
        // 2) Velocity smoothness
        for (int i = 1; i <= N - 2; ++i)
        {
            problem.AddResidualBlock(new ceres::AutoDiffCostFunction<VelSmoothResidual, 1, 4, 4, 4>(
                                         new VelSmoothResidual(opts.w_vel_smooth)),
                                     nullptr,
                                     &traj[i - 1].x,
                                     &traj[i].x,
                                     &traj[i + 1].x);
        }
 
        // 3) Optional jerk residuals
        if (opts.use_jerk && N >= 4)
        {
            for (int i = 2; i <= N - 2; ++i)
            {
                problem.AddResidualBlock(
                    new ceres::AutoDiffCostFunction<PoseJerkResidual, 3, 4, 4, 4, 4>(
                        new PoseJerkResidual(opts.w_pose_jerk)),
                    nullptr,
                    &traj[i - 2].x,
                    &traj[i - 1].x,
                    &traj[i].x,
                    &traj[i + 1].x);
            }
        }
 
        // 4) Robot smoothness residuals (optional)
        if (opts.use_robot_smooth)
        {
            for (int i = 1; i <= N - 2; ++i)
            {
                problem.AddResidualBlock(
                    new ceres::AutoDiffCostFunction<RobotSmoothResidual, 3, 4, 4, 4>(
                        new RobotSmoothResidual(opts.w_robot_smooth, -0.5, 0.0)),
                    nullptr,
                    &traj[i - 1].x,
                    &traj[i].x,
                    &traj[i + 1].x);
            }
        }
 
        // 5) Spacing residuals to encourage equal spacing
        double d_avg = computeAverageSpacing(traj);
        if (d_avg <= 1e-9)
            d_avg = 0.1;
        for (int i = 0; i < N - 1; ++i)
        {
            problem.AddResidualBlock(new ceres::AutoDiffCostFunction<SpacingResidual, 1, 4, 4>(
                                         new SpacingResidual(d_avg, opts.w_spacing)),
                                     nullptr,
                                     &traj[i].x,
                                     &traj[i + 1].x);
        }
 
        // 6) Obstacle residuals with arc-length weighting
        if (opts.use_obs)
        {
            for (int i = 0; i < N; ++i)
            {
                double s_local;
                if (i == 0)
                    s_local = segment_length(traj[0], traj[1]);
                else if (i == N - 1)
                    s_local = segment_length(traj[N - 2], traj[N - 1]);
                else
                    s_local = 0.5 * (segment_length(traj[i], traj[i - 1]) +
                                     segment_length(traj[i + 1], traj[i]));
                double s_sqrt = std::sqrt(std::max(1e-6, s_local));
                // problem.AddResidualBlock(new ceres::AutoDiffCostFunction<ObsResidual, 1, 4>(
                //                              new ObsResidual(grid, opts.clearance, s_sqrt, opts.w_obs)),
                //                          nullptr,
                //                          &traj[i].x);
                //
 
                // NOTE: Arc-length weighting is currently not working (not updated at each step)
                // Therefore, we don't even pass s_local to the ObsResidual
 
                problem.AddResidualBlock(new ceres::AutoDiffCostFunction<ObstacleResidual, 1, 4>(
                                             new ObstacleResidual(costmap_wrapper,
                                                                  opts.w_obs,
                                                                  opts.obsMaxDistance,
                                                                  opts.obsCollisionResidual)),
                                         nullptr,
                                         &traj[i].x);
            }
        }
 
        // 7) Velocity limit hinge residuals
        for (int i = 0; i < N; ++i)
        {
            problem.AddResidualBlock(new ceres::AutoDiffCostFunction<VelLimitResidual, 1, 4>(
                                         new VelLimitResidual(opts.vmax, opts.w_vel_limit)),
                                     nullptr,
                                     &traj[i].x);
        }
 
        // 8) Tracking residuals to original (optional)
        if (opts.use_tracking)
        {
            if (static_cast<int>(traj_original.size()) != N)
            {
                std::cerr << "[optimizeTrajectoryCeres] traj_original size mismatch; skipping "
                             "tracking.\n";
            }
            else
            {
                for (int i = 0; i < N; ++i)
                {
                    double wx = opts.w_track;
                    double wy = opts.w_track;
                    double wth = opts.w_track * 300.0;
                    double wv = opts.w_track * 0.2;
                    problem.AddResidualBlock(
                        new ceres::AutoDiffCostFunction<TrackingResidual, 4, 4>(
                            new TrackingResidual(traj_original[i].x,
                                                 traj_original[i].y,
                                                 traj_original[i].theta,
                                                 traj_original[i].v,
                                                 wx,
                                                 wy,
                                                 wth,
                                                 wv)),
                        nullptr,
                        &traj[i].x);
                }
            }
        }
 
        // Solver options: limit iterations (we recommend a small number for real-time)
        ceres::Solver::Options options;
        options.max_num_iterations = opts.max_iterations;
        options.linear_solver_type = ceres::SPARSE_NORMAL_CHOLESKY;
        options.minimizer_progress_to_stdout = false;
        options.num_threads = opts.num_threads;
        options.function_tolerance = 1e-6;
        options.parameter_tolerance = 1e-8;
 
 
        // Use Levenberg-Marquardt (better for constrained problems)
        options.trust_region_strategy_type = ceres::LEVENBERG_MARQUARDT;
        // Start with small trust region
        options.initial_trust_region_radius = 1e1; // Tune this: 1-100
        options.max_trust_region_radius = 1e3;
        options.min_trust_region_radius = 1e-3;
 
        // Conservative updates
        options.min_relative_decrease = 1e-3; // Require actual improvement
 
        auto print_residuals = [&problem, N]()
        {
            std::vector<ceres::ResidualBlockId> to_eval;
            problem.GetResidualBlocks(&to_eval);
            ceres::Problem::EvaluateOptions options;
            options.residual_blocks = to_eval;
            double total_cost = 0.0;
            std::vector<double> evaluated_residuals;
            problem.Evaluate(options, &total_cost, &evaluated_residuals, nullptr, nullptr);
            for (auto i = 0; i < evaluated_residuals.size(); i++)
            {
                ARMARX_DEBUG << i << ": " << evaluated_residuals[i];
            }
        };
 
        ceres::Solver::Summary summary;
        print_residuals();
        ceres::Solve(options, &problem, &summary);
        print_residuals();
        // Normalize angles after optimization (wrapping to [-pi,pi])
        normalizeTrajectoryAngles(traj);
 
        ARMARX_INFO << summary.FullReport();
 
        ARMARX_INFO << "[optimizeTrajectoryCeres] Done. Final cost: " << summary.final_cost
                    << " iterations: " << summary.num_successful_steps << "\n";
    }
 
} // namespace armarx::navigation::algorithms::orientation_aware::smoothing