external_wrench_plugin.hpp

// Copyright 2026 PAL Robotics S.L.
//
// 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 MUJOCO_ROS2_CONTROL_PLUGINS__EXTERNAL_WRENCH_PLUGIN_HPP_
#define MUJOCO_ROS2_CONTROL_PLUGINS__EXTERNAL_WRENCH_PLUGIN_HPP_
 
#include <mutex>
#include <queue>
#include <string>
#include <vector>
 
#include <mujoco_ros2_control_msgs/srv/apply_external_wrench.hpp>
#include <rclcpp/rclcpp.hpp>
#include <realtime_tools/realtime_publisher.hpp>
#include <visualization_msgs/msg/marker_array.hpp>
 
#include "mujoco_ros2_control_plugins/mujoco_ros2_control_plugins_base.hpp"
 
namespace mujoco_ros2_control_plugins
{
 
class ExternalWrenchPlugin : public MuJoCoROS2ControlPluginBase
{
public:
  ExternalWrenchPlugin() = default;
  ~ExternalWrenchPlugin() override = default;
 
  bool init(rclcpp::Node::SharedPtr node, const mjModel* model, mjData* data) override;
  void pre_step(mjData* data) override;
  void cleanup() override;
 
  void publish_markers(visualization_msgs::msg::MarkerArray& markers) const;
 
private:
  using ApplyExternalWrench = mujoco_ros2_control_msgs::srv::ApplyExternalWrench;
  using MarkerArray = visualization_msgs::msg::MarkerArray;
 
  struct ActiveWrench
  {
    int body_id{ -1 };
    std::string link_name;
    mjtNum force[3]{ 0.0, 0.0, 0.0 };              
    mjtNum torque[3]{ 0.0, 0.0, 0.0 };             
    mjtNum application_point[3]{ 0.0, 0.0, 0.0 };  
    rclcpp::Time end_time{ 0, 0, RCL_ROS_TIME };
    rclcpp::Duration ramp_down_duration{ 0, 0 };
    double current_scale{ 1.0 };  
  };
 
  void handleApplyWrench(const ApplyExternalWrench::Request::SharedPtr request,
                         ApplyExternalWrench::Response::SharedPtr response);
 
  void publish_markers();
 
  // ROS interfaces
  rclcpp::Node::SharedPtr node_;
  rclcpp::Logger logger_{ rclcpp::get_logger("ExternalWrenchPlugin") };
  rclcpp::Service<ApplyExternalWrench>::SharedPtr service_;
 
  rclcpp::Publisher<MarkerArray>::SharedPtr marker_pub_raw_;
  std::unique_ptr<realtime_tools::RealtimePublisher<MarkerArray>> marker_pub_;
 
  // Model pointer (const, valid for simulation lifetime)
  const mjModel* model_{ nullptr };
 
  // Pending queue written by service callback, drained in pre_step()
  std::mutex pending_mutex_;
  std::queue<ActiveWrench> pending_wrenches_;
 
  // Active wrenches — only modified inside pre_step()
  std::vector<ActiveWrench> active_wrenches_;
 
  // Body IDs whose xfrc_applied slots were written in the previous pre_step() call.
  // Zeroed at the start of the next pre_step() before re-accumulating.
  std::vector<int> prev_written_body_ids_;
 
  // Marker visualization scaling
  double force_arrow_scale_{ 0.01 };
  double torque_arrow_scale_{ 0.1 };
 
  std::atomic_bool service_requested_{ false };
};
 
}  // namespace mujoco_ros2_control_plugins
 
#endif  // MUJOCO_ROS2_CONTROL_PLUGINS__EXTERNAL_WRENCH_PLUGIN_HPP_