Program Listing for File differential_drive_model.hpp

↰ Return to documentation for file (include/beluga/motion/differential_drive_model.hpp)

// Copyright 2022-2023 Ekumen, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef BELUGA_MOTION_DIFFERENTIAL_DRIVE_MODEL_HPP
#define BELUGA_MOTION_DIFFERENTIAL_DRIVE_MODEL_HPP

#include <random>
#include <sophus/se3.hpp>
#include <tuple>

#include <beluga/type_traits/tuple_traits.hpp>

#include <beluga/3d_embedding.hpp>
#include <sophus/se2.hpp>
#include <sophus/so2.hpp>
#include <type_traits>

namespace beluga {


struct DifferentialDriveModelParam {

  double rotation_noise_from_rotation;

  double rotation_noise_from_translation;

  double translation_noise_from_translation;

  double translation_noise_from_rotation;

  double distance_threshold = 0.01;
};


template <class StateType = Sophus::SE2d>
class DifferentialDriveModel {
  static_assert(
      std::is_same_v<StateType, Sophus::SE2d> or std::is_same_v<StateType, Sophus::SE3d>,
      "Differential model only supports SE2 and SE3 state types.");

 public:
  using state_type = StateType;

  using control_type = std::tuple<state_type, state_type>;

  using param_type = DifferentialDriveModelParam;


  explicit DifferentialDriveModel(const param_type& params) : params_{params} {}


  template <class Control, typename = common_tuple_type_t<Control, control_type>>
  [[nodiscard]] auto operator()(const Control& action) const {
    const auto& [pose, previous_pose] = action;
    if constexpr (std::is_same_v<state_type, Sophus::SE2d>) {
      return sampling_fn_2d(pose, previous_pose);
    } else {
      return sampling_fn_3d(pose, previous_pose);
    }
  }

 private:
  using control_type_2d = std::tuple<Sophus::SE2d, Sophus::SE2d>;
  using control_type_3d = std::tuple<Sophus::SE3d, Sophus::SE3d>;

  [[nodiscard]] auto sampling_fn_3d(const Sophus::SE3d& pose, const Sophus::SE3d& previous_pose) const {
    const auto current_pose_2d = To2d(pose);
    const auto previous_pose_pose_2d = To2d(previous_pose);
    const auto two_d_sampling_fn = sampling_fn_2d(current_pose_2d, previous_pose_pose_2d);
    return [=](const state_type& state, auto& gen) { return To3d(two_d_sampling_fn(To2d(state), gen)); };
  }

  [[nodiscard]] auto sampling_fn_2d(const Sophus::SE2d& pose, const Sophus::SE2d& previous_pose) const {
    const auto translation = pose.translation() - previous_pose.translation();
    const double distance = translation.norm();
    const double distance_variance = distance * distance;

    const auto& previous_orientation = previous_pose.so2();
    const auto& current_orientation = pose.so2();
    const auto heading_rotation = Sophus::SO2d{std::atan2(translation.y(), translation.x())};
    const auto first_rotation =
        distance > params_.distance_threshold ? heading_rotation * previous_orientation.inverse() : Sophus::SO2d{};
    const auto second_rotation = current_orientation * previous_orientation.inverse() * first_rotation.inverse();

    using DistributionParam = typename std::normal_distribution<double>::param_type;
    const auto first_rotation_params = DistributionParam{
        first_rotation.log(), std::sqrt(
                                  params_.rotation_noise_from_rotation * rotation_variance(first_rotation) +
                                  params_.rotation_noise_from_translation * distance_variance)};
    const auto translation_params = DistributionParam{
        distance, std::sqrt(
                      params_.translation_noise_from_translation * distance_variance +
                      params_.translation_noise_from_rotation *
                          (rotation_variance(first_rotation) + rotation_variance(second_rotation)))};
    const auto second_rotation_params = DistributionParam{
        second_rotation.log(), std::sqrt(
                                   params_.rotation_noise_from_rotation * rotation_variance(second_rotation) +
                                   params_.rotation_noise_from_translation * distance_variance)};

    return [=](const auto& state, auto& gen) {
      static thread_local auto distribution = std::normal_distribution<double>{};
      const auto first_rotation = Sophus::SO2d{distribution(gen, first_rotation_params)};
      const auto translation = Eigen::Vector2d{distribution(gen, translation_params), 0.0};
      const auto second_rotation = Sophus::SO2d{distribution(gen, second_rotation_params)};
      return state * Sophus::SE2d{first_rotation, Eigen::Vector2d{0.0, 0.0}} *
             Sophus::SE2d{second_rotation, translation};
    };
  }
  param_type params_;

  static double rotation_variance(const Sophus::SO2d& rotation) {
    // Treat backward and forward motion symmetrically for the noise models.
    static const auto kFlippingRotation = Sophus::SO2d{Sophus::Constants<double>::pi()};
    const auto flipped_rotation = rotation * kFlippingRotation;
    const auto delta = std::min(std::abs(rotation.log()), std::abs(flipped_rotation.log()));
    return delta * delta;
  }
};

using DifferentialDriveModel2d = DifferentialDriveModel<Sophus::SE2d>;

using DifferentialDriveModel3d = DifferentialDriveModel<Sophus::SE3d>;

}  // namespace beluga

#endif