Costmap3DWrapper.h

Go to the documentation of this file.

#pragma once
#include <algorithm>
#include <cmath>
#include <functional>
#include <optional>
 
#include <ArmarXCore/core/logging/Logging.h>
 
#include <armarx/navigation/algorithms/orientation_aware/Costmap3D.h>
 
#include <ceres/ceres.h>
 
namespace armarx::navigation::algorithms::orientation_aware::smoothing
{
 
    //
    // A simple differentiable wrapper around Costmap3D.
    // Produces a *continuous* obstacle cost for Ceres AutoDiff.
    //
    // Design choices:
    //  * bilinear interpolation in X/Y
    //  * nearest-neighbor in orientation (your grid is discrete anyway)
    //  * returns 0 if value() returns nullopt
    //  * everything templated so Ceres autodiff works
    //
    // Generated by GPT-5 on 2025-12-04
 
    class Costmap3DWrapper
    {
    public:
        using Position = Eigen::Vector2f;
        using RotationDegrees = float;
        using Index = Eigen::Array2i;
 
    private:
        // self-written functions
        template <typename IndexT, typename WeightT, int size>
        struct LinearInterpolationResult
        {
            std::array<IndexT, size> indices;
            std::array<WeightT, size> weights;
 
            template <typename T>
            T
            calcWeightedAccess(std::function<T(const IndexT&)> accessor) const
            {
                T result = T(0);
                for (int i = 0; i < size; ++i)
                {
                    result += weights[i] * accessor(indices[i]);
                }
                return result;
            }
        };
 
        template <typename IndexT1, typename IndexT2, typename CommonWeightT, int size1, int size2>
        static LinearInterpolationResult<std::tuple<IndexT1, IndexT2>, CommonWeightT, size1 * size2>
        combineLinearInterpolationResults(
            LinearInterpolationResult<IndexT1, CommonWeightT, size1> r1,
            LinearInterpolationResult<IndexT2, CommonWeightT, size2> r2)
        {
            LinearInterpolationResult<std::tuple<IndexT1, IndexT2>, CommonWeightT, size1 * size2>
                result;
            for (int i = 0; i < size1; ++i)
            {
                for (int j = 0; j < size2; ++j)
                {
                    result.indices[i * size2 + j] = std::make_tuple(r1.indices[i], r2.indices[j]);
                    result.weights[i * size2 + j] = r1.weights[i] * r2.weights[j];
                }
            }
            return result;
        }
 
    public:
        Costmap3DWrapper(const armarx::navigation::algorithms::orientation_aware::Costmap3D& map) :
            map_(map)
        {
            cell_size_ = map_.params().cellSize; // [mm]
            inv_cell_size_ = 1.0f / cell_size_;
            grid_size_ = map_.getSize();
            map_.convertToPseudoSDF();
        }
 
        /// Templated call operator – fully autodiff-friendly
        template <typename T>
        bool
        operator()(const T* const x_ptr,
                   const T* const y_ptr,
                   const T* const theta_ptr,
                   T* out_cost) const
        {
            // Idea: First get indices for accessing using provided methods
            //       Then calculate interpolation using Jet types
 
            const auto rotation_indices = closestRotationIndexFromDegreesBilinear(*theta_ptr);
            const auto position_indices = toIndexBilinear(Eigen::Matrix<T, 2, 1>{*x_ptr, *y_ptr});
            const auto combined_interpolation =
                combineLinearInterpolationResults(rotation_indices, position_indices);
            out_cost[0] = combined_interpolation.template calcWeightedAccess<T>(
                [this](const std::tuple<int, Eigen::Array<int, 2, 1, 0, 2, 1>>& idx) -> T
                {
                    auto& [rot, pos] = idx;
                    Costmap3D::Index posIdx = {pos.x(), pos.y()};
                    auto v = map_.value(posIdx, rot);
 
                    if (!v.has_value())
                    {
                        // Masked or out-of-bounds cell → treat as deep obstacle so the
                        // optimizer (and hill-climber) stay away from unknown regions.
                        return T(-100.0);
                    }
 
                    T value = T(v.value());
 
                    // NOTE: We do NOT subtract an artificial safety margin here.
                    // The Costmap3D already stores the true signed distance of the full
                    // robot footprint to obstacles (with obstacleSafetyMargin baked in
                    // during construction). Adding another margin here would distort the
                    // distance field and break the collision gradient.
                    return value;
                });
            return true;
        }
 
        /// Query the costmap at a single discrete orientation (no theta interpolation).
        /// Position is still bilinearly interpolated. Differentiable w.r.t. (x,y).
        template <typename T>
        bool
        queryDiscreteOrientation(const T* const x_ptr,
                                 const T* const y_ptr,
                                 int rot_index,
                                 T* out_cost) const
        {
            const auto position_indices = toIndexBilinear(Eigen::Matrix<T, 2, 1>{*x_ptr, *y_ptr});
 
            LinearInterpolationResult<Costmap3D::RotationIndex, T, 1> rot_result;
            rot_result.indices[0] = wrapRotationIndex(rot_index);
            rot_result.weights[0] = T(1.0);
 
            const auto combined_interpolation =
                combineLinearInterpolationResults(rot_result, position_indices);
            out_cost[0] = combined_interpolation.template calcWeightedAccess<T>(
                [this](const std::tuple<int, Eigen::Array<int, 2, 1, 0, 2, 1>>& idx) -> T
                {
                    auto& [rot, pos] = idx;
                    Costmap3D::Index posIdx = {pos.x(), pos.y()};
                    auto v = map_.value(posIdx, rot);
                    if (!v.has_value())
                        return T(-100.0);
                    return T(v.value());
                });
            return true;
        }
 
    private:
        armarx::navigation::algorithms::orientation_aware::Costmap3D map_;
 
        float cell_size_;
        float inv_cell_size_;
        armarx::navigation::algorithms::orientation_aware::Costmap3D::Index grid_size_;
 
        template <typename T>
        static float
        toFloat(const T& val)
        {
            if constexpr (std::is_same_v<T, double> || std::is_same_v<T, float>)
                return static_cast<float>(val);
            else
                return static_cast<float>(val.a); // Jet
        }
 
        // ----------- helper: wrap [0,360) --------------
        template <typename T>
        static T
        wrapDegrees(const T& deg)
        {
            T d = deg;
            d = d - T(360.0) * ceres::floor(d / T(360.0));
            if (d > T(360.0))
                d -= T(360.0);
            if (d < T(0.0))
                d += T(360.0);
            return d;
        }
 
        int
        wrapRotationIndex(const int index) const
        {
            int d = index;
            if (d > int(map_.parameters.orientations - 1))
                d -= int(map_.parameters.orientations);
            if (d < int(0.0))
                d += int(map_.parameters.orientations);
            return d;
        }
 
        template <typename T>
        LinearInterpolationResult<Costmap3D::RotationIndex, T, 2>
        closestRotationIndexFromDegreesBilinear(const T& degrees_in) const
        {
            const T degrees = wrapDegrees(degrees_in);
            const T deg_per_orientation = T(360.F) / T(map_.parameters.orientations);
 
            // simply floor for lower one
            const int lower_index = wrapRotationIndex(
                static_cast<int>(toFloat(ceres::floor(degrees / deg_per_orientation))));
            const int higher_index = wrapRotationIndex(lower_index + 1);
 
            if (lower_index < 0 || lower_index > map_.parameters.orientations)
            {
                ARMARX_DEBUG << "lower_index out of range: " << lower_index;
                ARMARX_DEBUG << "degrees: " << degrees;
                ARMARX_DEBUG << VAROUT(deg_per_orientation);
            }
            if (lower_index < 0 || higher_index > map_.parameters.orientations)
            {
                ARMARX_DEBUG << "higher_index out of range: " << higher_index;
            }
 
            const T error_to_lower_normalized =
                (degrees - T(lower_index) * deg_per_orientation) / deg_per_orientation;
 
            return {.indices = {lower_index, higher_index},
                    .weights = {T(1) - error_to_lower_normalized, error_to_lower_normalized}};
        }
 
        template <typename T>
        static Eigen::Transform<T, 2, 1>
        convertTransform(const Eigen::Transform<float, 2, 1>& tf)
        {
            Eigen::Transform<T, 2, 1> out;
            out.matrix() = tf.matrix().template cast<T>();
            return out;
        }
 
        template <typename T>
        LinearInterpolationResult<Costmap3D::Index, T, 4>
        toIndexBilinear(const Eigen::Matrix<T, 2, 1>& globalPosition) const
        {
            // Transform world → local
            auto tfT = convertTransform<T>(map_.global_T_Costmap3D);
            Eigen::Matrix<T, 2, 1> localPosition = tfT.inverse() * globalPosition;
 
            const T cell = T(map_.parameters.cellSize);
 
            // Compute continuous cell coordinates
            const T vX = (localPosition.x() - cell / T(2) - T(map_.sceneBounds.min.x())) / cell;
 
            const T vY = (localPosition.y() - cell / T(2) - T(map_.sceneBounds.min.y())) / cell;
 
            // Bilinear kernel: continuous cell coordinates must be bracketed
            const T vX_shift = vX - T(0.01);
            const T vY_shift = vY - T(0.01);
 
            // Jet-safe floor/ceil
            int iXlow = static_cast<int>(toFloat(ceres::floor(vX_shift)));
            int iXhigh = static_cast<int>(toFloat(ceres::ceil(vX_shift)));
            int iYlow = static_cast<int>(toFloat(ceres::floor(vY_shift)));
            int iYhigh = static_cast<int>(toFloat(ceres::ceil(vY_shift)));
 
            // Clamp to map bounds
            const auto size = map_.getSize();
            int iXlow_cl = std::clamp(iXlow, 0, size.x() - 1);
            int iXhigh_cl = std::clamp(iXhigh, 0, size.x() - 1);
            int iYlow_cl = std::clamp(iYlow, 0, size.y() - 1);
            int iYhigh_cl = std::clamp(iYhigh, 0, size.y() - 1);
 
            // Bilinear weights (fully Jet-differentiable).
            // Note: we intentionally use the UNCLAMPED indices for weights so that
            // w00 + w10 + w01 + w11 == 1 even when the point is outside the grid
            // and the clamped indices collapse to the same cell.  This gives
            // correct nearest-neighbour extrapolation at the boundary.
            const T w00 = (T(iXhigh) - vX) * (T(iYhigh) - vY); // (low, low)
            const T w10 = (vX - T(iXlow)) * (T(iYhigh) - vY); // (high, low)
            const T w01 = (T(iXhigh) - vX) * (vY - T(iYlow)); // (low, high)
            const T w11 = (vX - T(iXlow)) * (vY - T(iYlow)); // (high, high)
 
            LinearInterpolationResult<Costmap3D::Index, T, 4> result;
 
            result.indices[0] = {iXlow_cl, iYlow_cl};
            result.indices[1] = {iXhigh_cl, iYlow_cl};
            result.indices[2] = {iXlow_cl, iYhigh_cl};
            result.indices[3] = {iXhigh_cl, iYhigh_cl};
 
            result.weights[0] = w00;
            result.weights[1] = w10;
            result.weights[2] = w01;
            result.weights[3] = w11;
 
            // Degeneracy guard: when the query point lands exactly on a grid line
            // (iXlow == iXhigh or iYlow == iYhigh), the bilinear weights sum to 0
            // and would incorrectly return 0 regardless of the actual cell value.
            // Fall back to nearest-neighbour in that case.
            T sumW = w00 + w10 + w01 + w11;
            if (ceres::abs(sumW) < T(1e-6))
            {
                T bestDist = T(std::numeric_limits<double>::max());
                int bestIdx = 0;
                for (int k = 0; k < 4; ++k)
                {
                    T ddx = T(result.indices[k].x()) - vX;
                    T ddy = T(result.indices[k].y()) - vY;
                    T d = ddx * ddx + ddy * ddy;
                    if (d < bestDist)
                    {
                        bestDist = d;
                        bestIdx = k;
                    }
                }
                for (int k = 0; k < 4; ++k)
                    result.weights[k] = T(0);
                result.weights[bestIdx] = T(1);
            }
 
            return result;
        }
    };
} // namespace armarx::navigation::algorithms::orientation_aware::smoothing