utils.hpp

 
#pragma once
 
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <hardware_interface/hardware_info.hpp>
#include <rclcpp/rclcpp.hpp>
 
#include "mujoco_ros2_control/data.hpp"
 
#include <cassert>
#include <cmath>
#include <random>
#include <string>
 
namespace mujoco_ros2_control
{
 
inline std::optional<hardware_interface::ComponentInfo>
get_sensor_from_info(const hardware_interface::HardwareInfo& hardware_info, const std::string& name)
{
  for (size_t sensor_index = 0; sensor_index < hardware_info.sensors.size(); sensor_index++)
  {
    const auto& sensor = hardware_info.sensors.at(sensor_index);
    if (hardware_info.sensors.at(sensor_index).name == name)
    {
      return sensor;
    }
  }
  return std::nullopt;
}
 
inline std::string get_sensor_parameter_or(const hardware_interface::ComponentInfo& sensor, const std::string& key,
                                           const std::string& default_value)
{
  if (auto it = sensor.parameters.find(key); it != sensor.parameters.end())
  {
    return it->second;
  }
  return default_value;
}
 
inline NoiseDistribution get_noise_distribution(const hardware_interface::ComponentInfo& sensor)
{
  const std::string value = get_sensor_parameter_or(sensor, "noise_distribution", "gaussian");
  if (value != "gaussian" && value != "uniform")
  {
    RCLCPP_WARN(rclcpp::get_logger("mujoco_ros2_control"),
                "Only noise distributions of 'gaussian' or 'uniform' are allowed, but you selected '%s'. "
                "Defaulting to 'gaussian'.",
                value.c_str());
  }
  return value == "uniform" ? NoiseDistribution::kUniform : NoiseDistribution::kGaussian;
}
 
inline void set_diagonal_covariance(std::vector<double>& covariance, double stddev)
{
  assert(covariance.size() == 9 && "covariance must be a flat 3x3 (9-element) matrix");
  covariance[0] = covariance[4] = covariance[8] = stddev * stddev;
}
 
namespace detail
{
 
template <int N>
inline Eigen::Matrix<double, N, 1> sample_noise(double stddev, NoiseDistribution distribution, std::mt19937& rng)
{
  Eigen::Matrix<double, N, 1> samples;
  if (distribution == NoiseDistribution::kUniform)
  {
    const double bound = stddev * std::sqrt(3.0);
    std::uniform_real_distribution<double> dist(-bound, bound);
    for (int i = 0; i < N; ++i)
    {
      samples[i] = dist(rng);
    }
  }
  else
  {
    std::normal_distribution<double> dist(0.0, stddev);
    for (int i = 0; i < N; ++i)
    {
      samples[i] = dist(rng);
    }
  }
  return samples;
}
 
}  // namespace detail
 
inline void add_sensor_noise(Eigen::Vector3d& value, double stddev, NoiseDistribution distribution, std::mt19937& rng)
{
  if (stddev <= 0.0)
  {
    return;
  }
  value += detail::sample_noise<3>(stddev, distribution, rng);
}
 
inline void add_sensor_noise(Eigen::Quaterniond& value, double stddev, NoiseDistribution distribution, std::mt19937& rng)
{
  if (stddev <= 0.0)
  {
    return;
  }
  value.coeffs() += detail::sample_noise<4>(stddev, distribution, rng);
  value.normalize();
}
 
inline void add_sensor_noise(Eigen::Vector3d& value, double stddev, NoiseState& noise)
{
  add_sensor_noise(value, stddev, noise.distribution, noise.rng);
}
 
inline void add_sensor_noise(Eigen::Quaterniond& value, double stddev, NoiseState& noise)
{
  add_sensor_noise(value, stddev, noise.distribution, noise.rng);
}
 
}  // namespace mujoco_ros2_control