mujoco_simulation.hpp

 
#pragma once
 
#include <algorithm>
#include <atomic>
#include <chrono>
#include <cmath>
#include <condition_variable>
#include <functional>
#include <memory>
#include <mutex>
#include <string>
#include <thread>
#include <vector>
 
#include <rclcpp/rclcpp.hpp>
#include <realtime_tools/realtime_publisher.hpp>
#include <rosgraph_msgs/msg/clock.hpp>
 
#include "glfw_adapter.h"  // for mj::GlfwAdapter
#include "simulate.h"      // must be on your include path, handled by CMake
 
#include <mujoco/mujoco.h>
 
#include <geometry_msgs/msg/pose_stamped.hpp>
#include <geometry_msgs/msg/twist_stamped.hpp>
#include <mujoco_ros2_control_msgs/msg/free_joint_state.hpp>
#include <mujoco_ros2_control_msgs/msg/simulation_state.hpp>
#include <mujoco_ros2_control_msgs/srv/reset_world.hpp>
#include <mujoco_ros2_control_msgs/srv/set_free_joint_state.hpp>
#include <mujoco_ros2_control_msgs/srv/set_pause.hpp>
#include <mujoco_ros2_control_msgs/srv/step_simulation.hpp>
#include <mujoco_ros2_control_plugins/mujoco_ros2_control_plugins_base.hpp>
#include <sensor_msgs/msg/joint_state.hpp>
 
namespace mujoco_ros2_control
{
 
namespace
{
 
constexpr double kIdleSleepFraction = 0.2;
 
}  // namespace
 
inline std::chrono::microseconds physics_loop_sleep_duration(double timestep)
{
  return std::min(std::chrono::microseconds(1000),
                  std::chrono::microseconds(static_cast<int64_t>(std::llround(timestep * kIdleSleepFraction * 1e6))));
}
 
class MujocoSimulation
{
public:
  using ResetCallback = std::function<void(bool fill_initial_state)>;
 
  using PreStepCallback = std::function<void(mjData* data)>;
 
  MujocoSimulation() = default;
 
  ~MujocoSimulation();
 
  bool initialize(rclcpp::Node::SharedPtr node, const std::string& model_path, const std::string& mujoco_model_topic,
                  double sim_speed_factor, bool headless);
 
  bool apply_keyframe(const std::string& keyframe_name);
 
  void capture_initial_state();
 
  void set_reset_callback(ResetCallback callback);
 
  void set_pre_step_callback(PreStepCallback callback);
 
  void start_physics_thread();
 
  void shutdown();
 
  mjModel* model()
  {
    return mj_model_;
  }
 
  mjData* data()
  {
    return mj_data_;
  }
 
  void reset_world_state(bool fill_initial_state);
 
  bool set_free_joint_states(const std::vector<mujoco_ros2_control_msgs::msg::FreeJointState>& free_joints,
                             std::string& error_message);
 
  void copy_physics_model(mjModel*& destination);
 
  void overwrite_physics_data(mjData* source);
 
  void copy_physics_data(mjData*& destination);
 
  mjData* acquire_data_snapshot();
 
  struct ControlState
  {
    mjtNum time{ 0.0 };
    std::vector<mjtNum> qpos;
    std::vector<mjtNum> qvel;
    std::vector<mjtNum> act;
    std::vector<mjtNum> qfrc_actuator;
    std::vector<mjtNum> sensordata;
    std::vector<mjtNum> ctrl;
  };
 
  void copy_control_state(ControlState& destination);
 
  void apply_control_data(mjData* control_data);
 
  std::recursive_mutex& mutex() const
  {
    RCLCPP_WARN_EXPRESSION(logger_, sim_mutex_ == nullptr, "Sim recursive mutex is still nullptr");
    return *sim_mutex_;
  }
 
  uint64_t step_count() const
  {
    return step_count_.load();
  }
 
private:
  void apply_staged_control_inputs();
 
  void publish_control_state();
 
  void refresh_data_snapshot();
 
  void physics_loop();
 
  void publish_clock();
 
  void update_sim_display();
 
  void reset_world_state(bool fill_initial_state, const mujoco_ros2_control_msgs::msg::SimulationState& state_overrides);
 
  int frame_body_id(const std::string& frame_id) const;
 
  bool validate_frame_id(const std::string& frame_id, const std::string& field_label, std::string& error_message) const;
 
  int find_free_joint_id(int body_id) const;
 
  bool validate_free_joint_states(const std::vector<mujoco_ros2_control_msgs::msg::FreeJointState>& free_joints,
                                  std::string& error_message);
 
  void apply_free_joint_states(const std::vector<mujoco_ros2_control_msgs::msg::FreeJointState>& free_joints);
 
  bool validate_joint_state_overrides(const sensor_msgs::msg::JointState& joint_state, std::string& error_message);
 
  void apply_joint_state_overrides(const sensor_msgs::msg::JointState& joint_state);
 
  // Service callbacks
  void reset_world_callback(const std::shared_ptr<mujoco_ros2_control_msgs::srv::ResetWorld::Request> request,
                            std::shared_ptr<mujoco_ros2_control_msgs::srv::ResetWorld::Response> response);
  void set_pause_callback(const std::shared_ptr<mujoco_ros2_control_msgs::srv::SetPause::Request> request,
                          std::shared_ptr<mujoco_ros2_control_msgs::srv::SetPause::Response> response);
  void step_simulation_callback(const std::shared_ptr<mujoco_ros2_control_msgs::srv::StepSimulation::Request> request,
                                std::shared_ptr<mujoco_ros2_control_msgs::srv::StepSimulation::Response> response);
  void
  set_free_joint_state_callback(const std::shared_ptr<mujoco_ros2_control_msgs::srv::SetFreeJointState::Request> request,
                                std::shared_ptr<mujoco_ros2_control_msgs::srv::SetFreeJointState::Response> response);
 
  rclcpp::Logger get_logger() const
  {
    return logger_;
  }
 
  // Logger
  rclcpp::Logger logger_ = rclcpp::get_logger("MujocoSimulation");
 
  // ROS node (owned by the HW interface, used here for services and clock publisher).
  rclcpp::Node::SharedPtr node_;
 
  // System information
  std::string model_path_;
  std::string mujoco_model_topic_;
 
  // MuJoCo data pointers, these are the primary containers used by the physics simulation.
  // It is generally not recommended to interact with them directly.
  mjModel* mj_model_{ nullptr };
 
  // Primary data container for the physics loop. We do not recommend interacting with this
  // directly unless you are sure of what you are doing.
  mjData* mj_data_{ nullptr };
 
  // Double-buffered snapshot of mj_data_: the physics loop pays for all scene-sized copies,
  // the consumer borrows the result via acquire_data_snapshot() and performs no copy at all.
  // The producer (any writer holding the sim mutex) fills snapshot_write_ outside any shared
  // lock and raises snapshot_ready_; the consumer, when the flag is up, swaps the two pointers
  // under data_exchange_mutex_ and reads snapshot_read_ in place.
  mjData* snapshot_write_{ nullptr };
  mjData* snapshot_read_{ nullptr };
  bool snapshot_ready_{ false };
 
  // Control inputs staged by apply_control_data, applied to mj_data_ before each step.
  std::vector<mjtNum> ctrl_staged_;
  std::vector<mjtNum> qfrc_applied_staged_;
 
  // False until apply_control_data is first called (and cleared on reset), so that initial /
  // reset ctrl values in mj_data_ are not clobbered by stale staging buffers.
  bool control_inputs_staged_{ false };
 
  // Guards only the snapshot pointer swap and snapshot_ready_ flag.
  // Lock order: sim_mutex_ (if needed) is always taken before this one.
  std::mutex data_exchange_mutex_;
 
  // Set by acquire_data_snapshot when a consumer takes the snapshot; the physics loop only
  // runs the expensive full refresh when this is set, so refresh bandwidth tracks consumer
  // demand instead of the batch rate. Starts true so the first refresh happens.
  std::atomic<bool> snapshot_refresh_requested_{ true };
 
  // Guards the staged control inputs (ctrl_staged_, qfrc_applied_staged_,
  // control_inputs_staged_). Separate from data_exchange_mutex_ so that staging commands in
  // write() and applying them before each physics step never queue behind a full mjData copy.
  // Critical sections are all small buffer copies.
  // Lock order: sim_mutex_ (if needed) before this one; never held with data_exchange_mutex_.
  std::mutex control_staging_mutex_;
 
  // Per-step control state served by copy_control_state. Guarded by its own mutex, separate
  // from data_exchange_mutex_, so the reduced control-state copies never queue behind a full
  // mjData snapshot copy.
  // Lock order: sim_mutex_ (if needed) before this one; never held with data_exchange_mutex_.
  ControlState control_state_;
  std::mutex control_state_mutex_;
 
  // For rendering
  mjvCamera cam_;
  mjvOption opt_;
  mjvPerturb pert_;
 
  // Speed scaling parameter. if set to >0 then we ignore the value set in the simulate app and instead
  // attempt to loop at whatever this is set to. If this is <0, then we use the value from the app.
  double sim_speed_factor_{ -1.0 };
 
  // True when running without a display (no GLFW window)
  bool headless_{ false };
 
  // Primary simulate object
  std::unique_ptr<mujoco::Simulate> sim_;
 
  // Threads for rendering physics and the UI simulation
  std::thread physics_thread_;
  std::thread ui_thread_;
  // Distinguishes a programmatic RenderLoop wakeup from the user closing its window.
  std::atomic<bool> explicit_shutdown_requested_{ false };
 
  // Primary clock publisher for the world
  std::shared_ptr<rclcpp::Publisher<rosgraph_msgs::msg::Clock>> clock_publisher_;
  realtime_tools::RealtimePublisher<rosgraph_msgs::msg::Clock>::SharedPtr clock_realtime_publisher_;
 
  // Mutex used inside simulate.h for protecting model/data, we keep a reference
  // here to protect access to shared data.
  // TODO: It would be far better to put all relevant data into a single container with accessors
  //       in a common location rather than passing around the raw pointer to the mutex, but it would
  //       require more work to pull it out of simulate.h.
  std::recursive_mutex* sim_mutex_{ nullptr };
 
  // Reset world service
  rclcpp::CallbackGroup::SharedPtr reset_world_cb_group_;
  rclcpp::Service<mujoco_ros2_control_msgs::srv::ResetWorld>::SharedPtr reset_world_service_;
 
  // Set pause service
  rclcpp::CallbackGroup::SharedPtr set_pause_cb_group_;
  rclcpp::Service<mujoco_ros2_control_msgs::srv::SetPause>::SharedPtr set_pause_service_;
 
  // Step simulation service
  rclcpp::CallbackGroup::SharedPtr step_simulation_cb_group_;
  rclcpp::Service<mujoco_ros2_control_msgs::srv::StepSimulation>::SharedPtr step_simulation_service_;
 
  // Set free joint state service (teleport/reset a free-joint object's pose)
  rclcpp::CallbackGroup::SharedPtr set_free_joint_state_cb_group_;
  rclcpp::Service<mujoco_ros2_control_msgs::srv::SetFreeJointState>::SharedPtr set_free_joint_state_service_;
 
  // Pending steps to execute while paused, and synchronization for blocking callers
  std::atomic<uint32_t> pending_steps_{ 0 };
  std::atomic<bool> step_diverged_{ false };
  std::atomic<bool> steps_interrupted_{ false };
  std::atomic<bool> keyboard_step_requested_{ false };
  std::atomic<uint64_t> step_count_{ 0 };
  std::mutex steps_cv_mutex_;
  std::condition_variable steps_cv_;
 
  // Storage for initial state (used for reset_world)
  std::vector<mjtNum> initial_qpos_;
  std::vector<mjtNum> initial_qvel_;
  std::vector<mjtNum> initial_ctrl_;
 
  // Callback into the HW interface to perform component-side reset bookkeeping.
  ResetCallback reset_callback_;
 
  // Callback run immediately before every mj_step(), defaults to nothing.
  PreStepCallback pre_step_callback_{ [](mjData* /*data*/) {} };
};
 
}  // namespace mujoco_ros2_control