Pause Functionality

mujoco_ros2_control/include/mujoco_ros2_control/mujoco_rendering.hpp

@@ -4,6 +4,7 @@
 #include "rclcpp/rclcpp.hpp"
 #include "mujoco/mujoco.h"
 #include "GLFW/glfw3.h"
+#include "mujoco_ros2_control/mujoco_ros2_control.hpp"
 
 namespace mujoco_ros2_control
 {
@@ -14,7 +15,7 @@ class MujocoRendering
   void operator=(const MujocoRendering &) = delete;
 
   static MujocoRendering* get_instance();
-  void init(rclcpp::Node::SharedPtr & node, mjModel* mujoco_model, mjData* mujoco_data);
+  void init(rclcpp::Node::SharedPtr & node, std::shared_ptr<MujocoRos2Control> & control, mjModel* mujoco_model, mjData* mujoco_data);
   bool is_close_flag_raised();
   void update();
   void close();
@@ -32,6 +33,7 @@ class MujocoRendering
 
   static MujocoRendering* instance_;
   rclcpp::Node::SharedPtr node_;  // TODO: delete node and add logger
+  std::shared_ptr<MujocoRos2Control> control_;
   mjModel* mj_model_;
   mjData* mj_data_;
   mjvCamera mjv_cam_;

mujoco_ros2_control/include/mujoco_ros2_control/mujoco_ros2_control.hpp

@@ -19,6 +19,9 @@ class MujocoRos2Control
   ~MujocoRos2Control();
   void init();
   void update();
+  void set_pause_simulation(bool pause);
+  bool is_paused() const;
+  void step_once();
 
 private:
   void publish_sim_time(rclcpp::Time sim_time);
@@ -37,6 +40,8 @@ class MujocoRos2Control
 
   rclcpp::Time last_update_sim_time_ros_;
   rclcpp::Publisher<rosgraph_msgs::msg::Clock>::SharedPtr clock_publisher_;
+
+  bool pause_simulation_;
 };
 }  // namespace mujoco_ros2_control
 

mujoco_ros2_control/src/mujoco_rendering.cpp

@@ -21,9 +21,12 @@ MujocoRendering::MujocoRendering()
 {
 }
 
-void MujocoRendering::init(rclcpp::Node::SharedPtr & node, mjModel* mujoco_model, mjData* mujoco_data)
+void MujocoRendering::init(rclcpp::Node::SharedPtr & node,
+                           std::shared_ptr<MujocoRos2Control> & control,
+                           mjModel* mujoco_model, mjData* mujoco_data)
 {
   node_ = node;
+  control_ = control;
   mj_model_ = mujoco_model;
   mj_data_ = mujoco_data;
 
@@ -117,6 +120,12 @@ void MujocoRendering::keyboard_callback_impl(GLFWwindow* window, int key, int sc
     mj_resetData(mj_model_, mj_data_);
     mj_forward(mj_model_, mj_data_);
   }
+  // space: pause/unpause simulation
+  if (act == GLFW_PRESS && key == GLFW_KEY_SPACE) {
+    if (control_) {
+      control_->set_pause_simulation(!control_->is_paused());
+    }
+  }
 }
 
 void MujocoRendering::mouse_button_callback_impl(GLFWwindow* window, int button, int act, int mods)

mujoco_ros2_control/src/mujoco_ros2_control.cpp

@@ -8,7 +8,7 @@ namespace mujoco_ros2_control
 {
 MujocoRos2Control::MujocoRos2Control(rclcpp::Node::SharedPtr & node, mjModel* mujoco_model, mjData* mujoco_data)
   : node_(node), mj_model_(mujoco_model), mj_data_(mujoco_data), logger_(rclcpp::get_logger(node_->get_name() + std::string(".mujoco_ros2_control"))),
-    control_period_(rclcpp::Duration(1, 0)), last_update_sim_time_ros_(0, 0, RCL_ROS_TIME)
+    control_period_(rclcpp::Duration(1, 0)), last_update_sim_time_ros_(0, 0, RCL_ROS_TIME), pause_simulation_(false)
 {
 }
 
@@ -137,19 +137,25 @@ void MujocoRos2Control::update()
   rclcpp::Duration sim_period = sim_time_ros - last_update_sim_time_ros_;
 
   publish_sim_time(sim_time_ros);
-
-  mj_step1(mj_model_, mj_data_);
-
+
+  if (!pause_simulation_) {
+    mj_step1(mj_model_, mj_data_);
+  }
+
   if (sim_period >= control_period_) {
     controller_manager_->read(sim_time_ros, sim_period);
     controller_manager_->update(sim_time_ros, sim_period);
-    last_update_sim_time_ros_ = sim_time_ros;
+    if (!pause_simulation_) {
+      last_update_sim_time_ros_ = sim_time_ros;
+    }
   }
-
+  
   // use same time as for read and update call - this is how it is done in ros2_control_node
   controller_manager_->write(sim_time_ros, sim_period);
-
-  mj_step2(mj_model_, mj_data_);
+
+  if (!pause_simulation_) {
+    mj_step2(mj_model_, mj_data_);
+  }
 }
 
 void MujocoRos2Control::publish_sim_time(rclcpp::Time sim_time)
@@ -160,4 +166,31 @@ void MujocoRos2Control::publish_sim_time(rclcpp::Time sim_time)
   clock_publisher_->publish(sim_time_msg);
 }
 
+void MujocoRos2Control::set_pause_simulation(bool pause)
+{
+  pause_simulation_ = pause;
+}
+
+bool MujocoRos2Control::is_paused() const
+{
+  return pause_simulation_;
+}
+
+void MujocoRos2Control::step_once()
+{
+  mj_step1(mj_model_, mj_data_);
+
+  controller_manager_->read(last_update_sim_time_ros_ + control_period_, control_period_);
+  controller_manager_->update(last_update_sim_time_ros_ + control_period_, control_period_);
+  controller_manager_->write(last_update_sim_time_ros_ + control_period_, control_period_);
+
+  mj_step2(mj_model_, mj_data_);
+
+  auto sim_time = mj_data_->time;
+  int sim_time_sec = static_cast<int>(sim_time);
+  int sim_time_nanosec = static_cast<int>((sim_time - sim_time_sec) * 1000000000);
+  last_update_sim_time_ros_ = rclcpp::Time(sim_time_sec, sim_time_nanosec, RCL_ROS_TIME);
+}
+
+
 } // namespace mujoco_ros2_control

mujoco_ros2_control/src/mujoco_ros2_control_node.cpp

@@ -17,6 +17,9 @@ int main(int argc, const char** argv) {
 
   RCLCPP_INFO_STREAM(node->get_logger(), "Initializing mujoco_ros2_control node...");
   auto model_path = node->get_parameter("mujoco_model_path").as_string();
+  bool pause_sim = node->has_parameter("pause_simulation")
+                      ? node->get_parameter("pause_simulation").as_bool()
+                      : node->declare_parameter("pause_simulation", false);
 
   // load and compile model
   char error[1000] = "Could not load binary model";
@@ -34,13 +37,19 @@ int main(int argc, const char** argv) {
   mujoco_data = mj_makeData(mujoco_model);
 
   // initialize mujoco control
-  auto control = mujoco_ros2_control::MujocoRos2Control(node, mujoco_model, mujoco_data);
-  control.init();
+  auto control = std::make_shared<mujoco_ros2_control::MujocoRos2Control>(node, mujoco_model, mujoco_data);
+  control->init();
+  control->set_pause_simulation(pause_sim);
   RCLCPP_INFO_STREAM(node->get_logger(), "Mujoco ros2 controller has been successfully initialized !");
 
-  // initialize mujoco redering
+  // if paused from start, step before rendering
+  if (control->is_paused()) {
+    control->step_once();
+  }
+
+  // initialize mujoco rendering
   auto rendering = mujoco_ros2_control::MujocoRendering::get_instance();
-  rendering->init(node, mujoco_model, mujoco_data);
+  rendering->init(node, control, mujoco_model, mujoco_data);
   RCLCPP_INFO_STREAM(node->get_logger(), "Mujoco rendering has been successfully initialized !");
 
   // run main loop, target real-time simulation and 60 fps rendering
@@ -49,9 +58,19 @@ int main(int argc, const char** argv) {
     //  Assuming MuJoCo can simulate faster than real-time, which it usually can,
     //  this loop will finish on time for the next frame to be rendered at 60 fps.
     //  Otherwise add a cpu timer and exit this loop when it is time to render.
-    mjtNum simstart = mujoco_data->time;
-    while (mujoco_data->time - simstart < 1.0/60.0) {
-      control.update();
+
+    if (control->is_paused()) {
+      // use wall-clock time
+      auto start_wall = std::chrono::steady_clock::now();
+      while (std::chrono::steady_clock::now() - start_wall < std::chrono::duration<double>(1.0 / 60.0)) {
+        control->update();
+      }
+    } else {
+      // simulation time
+      mjtNum simstart = mujoco_data->time;
+      while (mujoco_data->time - simstart < 1.0 / 60.0) {
+        control->update();
+      }
     }
     rendering->update();
   }