feat(sim): Interactive sim viewer in separate process

pkrack · web-flow · commit 65dfbb6931d5 · 2024-10-02T10:13:02.000+02:00
cf. PR #137 BREAKING CHANGE: removed `on_screen` sim camera
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -67,6 +67,29 @@ FetchContent_Declare(pybind11
     GIT_PROGRESS TRUE
     EXCLUDE_FROM_ALL
 )
+set(MuJoCo_UI_URL https://raw.githubusercontent.com/google-deepmind/mujoco/${MuJoCo_VERSION}/simulate/)
+set(MuJoCo_UI_FILES platform_ui_adapter.cc platform_ui_adapter.h glfw_adapter.cc glfw_adapter.h glfw_adapter.cc glfw_corevideo.h glfw_corevideo.mm glfw_dispatch.cc glfw_dispatch.h)
+# Set the destination directory for the downloaded files
+set(MuJoCo_UI_DIR "${CMAKE_BINARY_DIR}/mujoco_ui")
+
+# Create the directory if it doesn't exist
+file(MAKE_DIRECTORY ${MuJoCo_UI_DIR})
+
+# Loop through each file and download it only if it doesn't already exist
+foreach(file ${MuJoCo_UI_FILES})
+  set(output_file "${MuJoCo_UI_DIR}/${file}")
+  if (NOT EXISTS ${output_file})
+    message(STATUS "Downloading ${file}...")
+    file(DOWNLOAD ${MuJoCo_UI_URL}${file} ${output_file})
+  else()
+    message(STATUS "File ${file} already exists, skipping download.")
+  endif()
+endforeach()
+
+# Add the include directories or link the files as needed
+add_library(mujoco_ui_lib ${MuJoCo_UI_DIR}/glfw_adapter.cc ${MuJoCo_UI_DIR}/platform_ui_adapter.cc ${MuJoCo_UI_DIR}/glfw_dispatch.cc)
+target_link_libraries(mujoco_ui_lib MuJoCo::MuJoCo glfw)
+target_include_directories(mujoco_ui_lib INTERFACE ${MuJoCo_UI_DIR})
 
 FetchContent_MakeAvailable(libfranka rl pybind11)
 
diff --git a/python/examples/env_cartesian_control.py b/python/examples/env_cartesian_control.py
@@ -53,6 +53,7 @@ def main():
                 camera_set_cfg=default_mujoco_cameraset_cfg(),
                 max_relative_movement=0.5,
             )
+            env_rel.get_wrapper_attr("sim").open_gui()
 
         env_rel.reset()
         print(env_rel.unwrapped.robot.get_cartesian_position())
diff --git a/python/examples/env_joint_control.py b/python/examples/env_joint_control.py
@@ -56,6 +56,7 @@ def main():
                 camera_set_cfg=default_mujoco_cameraset_cfg(),
                 max_relative_movement=np.deg2rad(5),
             )
+            env_rel.get_wrapper_attr("sim").open_gui()
 
         for _ in range(10):
             obs, info = env_rel.reset()
diff --git a/python/examples/fr3.py b/python/examples/fr3.py
@@ -59,15 +59,16 @@ def main():
             # add camera to have a rendering gui
             cameras = {
                 "default_free": SimCameraConfig(
-                    identifier="", type=int(CameraType.default_free), on_screen_render=True
+                    identifier="", type=int(CameraType.default_free)
                 ),
                 "wrist": SimCameraConfig(
-                    identifier="eye-in-hand_0", type=int(CameraType.fixed), on_screen_render=False
+                    identifier="eye-in-hand_0", type=int(CameraType.fixed)
                 ),
                 # TODO: odd behavior when not both cameras are used: only last image is shown
             }
             cam_cfg = SimCameraSetConfig(cameras=cameras, resolution_width=1280, resolution_height=720, frame_rate=20)
             camera_set = SimCameraSet(simulation, cam_cfg)
+            simulation.open_gui()
 
         else:
             robot = rcsss.hw.FR3(ROBOT_IP, str(rcsss.scenes["lab"].parent / "fr3.urdf"))
diff --git a/python/rcsss/_core/sim.pyi b/python/rcsss/_core/sim.pyi
@@ -25,6 +25,7 @@ __all__ = [
     "default_free",
     "fixed",
     "free",
+    "open_gui_window",
     "tracking",
 ]
 M = typing.TypeVar("M", bound=int)
@@ -133,14 +134,15 @@ class FrankaHand(rcsss._core.common.Gripper):
 
 class Sim:
     def __init__(self, mjmdl: int, mjdata: int) -> None: ...
+    def _start_gui_server(self, id: str) -> None: ...
+    def _stop_gui_server(self) -> None: ...
     def is_converged(self) -> bool: ...
     def reset(self) -> None: ...
     def step(self, k: int) -> None: ...
     def step_until_convergence(self) -> None: ...
 
 class SimCameraConfig:
     identifier: str
-    on_screen_render: bool
     type: CameraType
     def __init__(self) -> None: ...
 
@@ -159,6 +161,8 @@ class SimCameraSetConfig:
     resolution_width: int
     def __init__(self) -> None: ...
 
+def open_gui_window(uuid: str) -> None: ...
+
 default_free: CameraType  # value = <CameraType.default_free: 3>
 fixed: CameraType  # value = <CameraType.fixed: 2>
 free: CameraType  # value = <CameraType.free: 0>
diff --git a/python/rcsss/camera/sim.py b/python/rcsss/camera/sim.py
@@ -20,7 +20,6 @@
 
 class SimCameraConfig(BaseCameraConfig):
     type: int  # CamBaseCameraConfigeraType
-    on_screen_render: bool
 
 
 class SimCameraSetConfig(BaseCameraSetConfig):
@@ -50,7 +49,6 @@ def get_type(t):
         for name, camera_cfg in cfg.cameras.items():
             cpp_camera_cfg = _SimCameraConfig()
             cpp_camera_cfg.type = get_type(camera_cfg.type)
-            cpp_camera_cfg.on_screen_render = camera_cfg.on_screen_render
             cpp_camera_cfg.identifier = camera_cfg.identifier
             cameras[name] = cpp_camera_cfg
 
diff --git a/python/rcsss/config.py b/python/rcsss/config.py
@@ -45,7 +45,7 @@ def create_sample_config_yaml(path: str):
         camera_config=camera_cfg,
         camera_type="realsense",
     )
-    cameras_sim = {"human_readable_name": SimCameraConfig(identifier="mjcf_name", type=0, on_screen_render=False)}
+    cameras_sim = {"human_readable_name": SimCameraConfig(identifier="mjcf_name", type=0)}
     sim = SimConfig(camera=SimCameraSetConfig(cameras=cameras_sim))
     cfg = Config(hw=hw, sim=sim)
     yml = to_yaml_str(cfg)
diff --git a/python/rcsss/control/liveplot.py b/python/rcsss/control/liveplot.py
@@ -2,6 +2,7 @@
 from itertools import product
 from socket import AF_INET, SOCK_DGRAM, socket
 from struct import unpack
+from typing import Optional
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -18,7 +19,7 @@
 titles = [["x", "rx"], ["y", "ry"], ["z", "rz"]]
 colors = [["r", "r"], ["b", "b"], ["g", "g"]]
 buffers: list[list[deque]] = [[deque(maxlen=buffsize) for j in range(2)] for i in range(3)]
-plots: list[list[plt.Figure]] = [[None for j in range(2)] for i in range(3)]
+plots: list[list[Optional[plt.Figure]]] = [[None for j in range(2)] for i in range(3)]
 
 fig, axs = plt.subplots(3, 2)
 for i, j in product(range(3), range(2)):
@@ -56,8 +57,9 @@
 
         for i, j in product(range(3), range(2)):
             buffer = buffers[i][j]
-            plots[i][j].set_xdata(np.arange(len(buffer)))
-            plots[i][j].set_ydata(buffer)
+            # mypy fails to coerce type after assert so we just ignore
+            plots[i][j].set_xdata(np.arange(len(buffer)))  # type: ignore
+            plots[i][j].set_ydata(buffer)  # type: ignore
             axs[i][j].relim()
             axs[i][j].autoscale_view()
             axs[i][j].set_title(titles[i][j] + ": {:.2f}".format(buffer[-1]))
diff --git a/python/rcsss/control/vive.py b/python/rcsss/control/vive.py
@@ -251,8 +251,8 @@ def sim():
     env_sim = FR3Sim(FR3Env(robot, ControlMode.JOINTS), simulation)
 
     cameras = {
-        "wrist": SimCameraConfig(identifier="eye-in-hand_0", type=int(CameraType.fixed), on_screen_render=False),
-        "default_free": SimCameraConfig(identifier="", type=int(CameraType.default_free), on_screen_render=True),
+        "wrist": SimCameraConfig(identifier="eye-in-hand_0", type=int(CameraType.fixed)),
+        "default_free": SimCameraConfig(identifier="", type=int(CameraType.default_free)),
     }
     cam_cfg = SimCameraSetConfig(cameras=cameras, resolution_width=640, resolution_height=480, frame_rate=10)
     camera_set = SimCameraSet(simulation, cam_cfg)
diff --git a/python/rcsss/envs/factories.py b/python/rcsss/envs/factories.py
@@ -114,8 +114,8 @@ def default_fr3_sim_gripper_cfg():
 def default_mujoco_cameraset_cfg():
 
     cameras = {
-        "wrist": SimCameraConfig(identifier="eye-in-hand_0", type=int(CameraType.fixed), on_screen_render=False),
-        "default_free": SimCameraConfig(identifier="", type=int(CameraType.default_free), on_screen_render=True),
+        "wrist": SimCameraConfig(identifier="eye-in-hand_0", type=int(CameraType.fixed)),
+        "default_free": SimCameraConfig(identifier="", type=int(CameraType.default_free)),
     }
     return SimCameraSetConfig(cameras=cameras, resolution_width=1280, resolution_height=720, frame_rate=10)
 
diff --git a/python/rcsss/envs/sim.py b/python/rcsss/envs/sim.py
@@ -142,12 +142,8 @@ def env_from_xml_paths(
             c_env = GripperWrapper(c_env, gripper)
         if camera:
             cameras = {
-                "wrist": SimCameraConfig(
-                    identifier="eye-in-hand_0", type=int(CameraType.fixed), on_screen_render=False
-                ),
-                "default_free": SimCameraConfig(
-                    identifier="", type=int(CameraType.default_free), on_screen_render=True
-                ),
+                "wrist": SimCameraConfig(identifier="eye-in-hand_0", type=int(CameraType.fixed)),
+                "default_free": SimCameraConfig(identifier="", type=int(CameraType.default_free)),
             }
             cam_cfg = SimCameraSetConfig(cameras=cameras, resolution_width=640, resolution_height=480, frame_rate=5)
             camera_set = SimCameraSet(simulation, cam_cfg)
diff --git a/python/rcsss/sim.py b/python/rcsss/sim.py
@@ -1,16 +1,24 @@
+import multiprocessing as mp
+import uuid
 from logging import getLogger
 from os import PathLike
 from pathlib import Path
+from typing import Optional
 
 import mujoco as mj
 from rcsss._core.sim import FR3, FHConfig, FHState, FR3Config, FR3State, FrankaHand
 from rcsss._core.sim import Sim as _Sim
+from rcsss._core.sim import open_gui_window as _open_gui_window
 
 __all__ = ["Sim", "FR3", "FR3Config", "FR3State", "FHConfig", "FHState", "FrankaHand"]
 
 logger = getLogger(__name__)
 
 
+def _start_gui(identifier: str):
+    _open_gui_window(identifier)
+
+
 class Sim(_Sim):
     def __init__(self, mjmdl: str | PathLike):
         mjmdl = Path(mjmdl)
@@ -23,3 +31,14 @@ def __init__(self, mjmdl: str | PathLike):
             logger.error(msg)
         self.data = mj.MjData(self.model)
         super().__init__(self.model._address, self.data._address)
+        self._gui_uuid: Optional[str] = None
+        self._mp_context = mp.get_context("spawn")
+
+    def open_gui(self):
+        if self._gui_uuid is None:
+            self._gui_uuid = "rcsss_" + str(uuid.uuid4())
+            self._start_gui_server(self._gui_uuid)
+        self._mp_context.Process(target=_start_gui, args=(self._gui_uuid,), daemon=True).start()
+
+    def __del__(self):
+        self._stop_gui_server()
diff --git a/src/pybind/rcsss.cpp b/src/pybind/rcsss.cpp
@@ -13,6 +13,7 @@
 #include <sim/FR3.h>
 #include <sim/FrankaHand.h>
 #include <sim/camera.h>
+#include <sim/gui.h>
 
 #include <memory>
 
@@ -438,7 +439,9 @@ PYBIND11_MODULE(_core, m) {
            py::call_guard<py::gil_scoped_release>())
       .def("is_converged", &rcs::sim::Sim::is_converged)
       .def("step", &rcs::sim::Sim::step, py::arg("k"))
-      .def("reset", &rcs::sim::Sim::reset);
+      .def("reset", &rcs::sim::Sim::reset)
+      .def("_start_gui_server", &rcs::sim::Sim::start_gui_server, py::arg("id"))
+      .def("_stop_gui_server", &rcs::sim::Sim::stop_gui_server);
   py::class_<rcs::sim::FrankaHand, rcs::common::Gripper,
              std::shared_ptr<rcs::sim::FrankaHand>>(sim, "FrankaHand")
       .def(py::init<std::shared_ptr<rcs::sim::Sim>, const std::string &,
@@ -465,9 +468,7 @@ PYBIND11_MODULE(_core, m) {
   py::class_<rcs::sim::SimCameraConfig>(sim, "SimCameraConfig")
       .def(py::init<>())
       .def_readwrite("identifier", &rcs::sim::SimCameraConfig::identifier)
-      .def_readwrite("type", &rcs::sim::SimCameraConfig::type)
-      .def_readwrite("on_screen_render",
-                     &rcs::sim::SimCameraConfig::on_screen_render);
+      .def_readwrite("type", &rcs::sim::SimCameraConfig::type);
   py::class_<rcs::sim::SimCameraSetConfig>(sim, "SimCameraSetConfig")
       .def(py::init<>())
       .def_readwrite("cameras", &rcs::sim::SimCameraSetConfig::cameras)
@@ -492,4 +493,5 @@ PYBIND11_MODULE(_core, m) {
       .def("get_latest_frameset", &rcs::sim::SimCameraSet::get_latest_frameset)
       .def("get_timestamp_frameset",
            &rcs::sim::SimCameraSet::get_timestamp_frameset, py::arg("ts"));
+  sim.def("open_gui_window", &rcs::sim::open_gui_window, py::arg("uuid"));
 }
diff --git a/src/sim/CMakeLists.txt b/src/sim/CMakeLists.txt
@@ -1,8 +1,8 @@
 add_library(sim)
 target_sources(sim
-  PRIVATE sim.cpp FR3.cpp renderer.cpp camera.cpp FrankaHand.cpp
+  PRIVATE sim.cpp FR3.cpp renderer.cpp camera.cpp FrankaHand.cpp gui_server.cpp gui_client.cpp
 )
-target_link_libraries(sim PUBLIC glfw common MuJoCo::MuJoCo mdl)
+target_link_libraries(sim PUBLIC glfw common MuJoCo::MuJoCo mdl mujoco_ui_lib)
 
 add_executable(test_sim test.cpp sim.cpp)
 target_link_libraries(test_sim PRIVATE sim MuJoCo::MuJoCo common mdl glfw)
diff --git a/src/sim/FR3.cpp b/src/sim/FR3.cpp
@@ -67,8 +67,6 @@ void FR3::construct() {
   this->rl.ik->setRandomRestarts(0);
   this->rl.ik->setEpsilon(1e-3);
   this->init_ids();
-  this->sim->register_cb(std::bind(&FR3::collision_callback, this),
-                         this->cfg.seconds_between_callbacks);
   this->sim->register_cb(std::bind(&FR3::is_arrived_callback, this),
                          this->cfg.seconds_between_callbacks);
   this->sim->register_cb(std::bind(&FR3::is_moving_callback, this),
diff --git a/src/sim/camera.cpp b/src/sim/camera.cpp
@@ -28,7 +28,7 @@ SimCameraSet::SimCameraSet(std::shared_ptr<Sim> sim, SimCameraSetConfig cfg)
         [this](const std::string& id, mjrContext& ctx, mjvScene& scene,
                mjvOption& opt) { this->frame_callback(id, ctx, scene, opt); },
         id, 1.0 / this->cfg.frame_rate, this->cfg.resolution_width,
-        this->cfg.resolution_height, !cam.on_screen_render);
+        this->cfg.resolution_height);
 
     mjvCamera mjcam;
     mjv_defaultCamera(&mjcam);
@@ -113,4 +113,4 @@ void SimCameraSet::frame_callback(const std::string& id, mjrContext& ctx,
 }
 
 }  // namespace sim
-}  // namespace rcs
+}  // namespace rcs
diff --git a/src/sim/camera.h b/src/sim/camera.h
@@ -24,7 +24,6 @@ enum CameraType {
 struct SimCameraConfig {
   std::string identifier;
   CameraType type;
-  bool on_screen_render;
 };
 struct SimCameraSetConfig {
   std::unordered_map<std::string, SimCameraConfig> cameras;
@@ -68,4 +67,4 @@ class SimCameraSet {
 };
 }  // namespace sim
 }  // namespace rcs
-#endif  // RCS_CAMSIM_H
+#endif  // RCS_CAMSIM_H
diff --git a/src/sim/gui.h b/src/sim/gui.h
@@ -0,0 +1,71 @@
+#ifndef GUI_H
+#define GUI_H
+#include <mujoco/mujoco.h>
+
+#include <bitset>
+#include <boost/interprocess/managed_shared_memory.hpp>
+#include <boost/interprocess/sync/named_sharable_mutex.hpp>
+#include <boost/interprocess/sync/named_upgradable_mutex.hpp>
+#include <optional>
+#include <string>
+
+namespace rcs {
+namespace sim {
+using namespace boost::interprocess;
+static const char* INFO_BYTE = "I";
+static const char* STATE = "S";
+static const char* MODEL = "M";
+static const std::string STATE_LOCK_POSTFIX = "_state_lock";
+static const std::string INFO_LOCK_POSTFIX = "_info_lock";
+static constexpr mjtState MJ_PHYSICS_SPEC = mjSTATE_FULLPHYSICS;
+
+static size_t calculate_state_size(mjModel const* m) {
+  return mj_stateSize(m, MJ_PHYSICS_SPEC);
+}
+
+static size_t calculate_mdl_size(mjModel const* m) { return mj_sizeModel(m); }
+
+static size_t calculate_shm_size(mjModel const* m, mjData const* d) {
+  size_t const extra_info_size = 1;
+  size_t const total_required_size = extra_info_size + calculate_mdl_size(m) +
+                                     calculate_state_size(m) * sizeof(mjtNum);
+  size_t const estimated_overhead =
+      total_required_size * 0.1;  // 10% overhead estimation
+  return total_required_size + estimated_overhead;
+}
+void open_gui_window(std::string& id);
+
+struct shm {
+  struct {
+    mjtNum* ptr = nullptr;
+    size_t size = 0;
+  } state;
+  struct {
+    char* ptr = nullptr;
+    size_t size = 0;
+  } model;
+  bool* info_byte = nullptr;  // extra info byte in shared memory
+  managed_shared_memory manager;
+  named_sharable_mutex state_lock;
+  named_upgradable_mutex info_lock;
+};
+
+class GuiServer {
+ public:
+  GuiServer(mjModel* m, mjData* d, const std::string& id);
+  ~GuiServer();
+  void update_mjdata_callback();
+
+ private:
+  // don't make this a reference, apparently references are not valid in
+  // destructors. or at least in this case the reference to id was not longer
+  // valid in the destructor.
+  const std::string id;
+  struct shm shm;
+  mjModel* m;
+  mjData* d;
+};
+
+}  // namespace sim
+}  // namespace rcs
+#endif
diff --git a/src/sim/gui_client.cpp b/src/sim/gui_client.cpp
diff --git a/src/sim/gui_server.cpp b/src/sim/gui_server.cpp
diff --git a/src/sim/renderer.cpp b/src/sim/renderer.cpp
diff --git a/src/sim/sim.cpp b/src/sim/sim.cpp
diff --git a/src/sim/sim.h b/src/sim/sim.h

Original file line number	Diff line number	Diff line change
`@@ -53,6 +53,7 @@ def main():`
`53`	`53`	`camera_set_cfg=default_mujoco_cameraset_cfg(),`
`54`	`54`	`max_relative_movement=0.5,`
`55`	`55`	`)`
	`56`	`+ env_rel.get_wrapper_attr("sim").open_gui()`
`56`	`57`
`57`	`58`	`env_rel.reset()`
`58`	`59`	`print(env_rel.unwrapped.robot.get_cartesian_position())`
Original file line number	Diff line number	Diff line change
`@@ -56,6 +56,7 @@ def main():`
`56`	`56`	`camera_set_cfg=default_mujoco_cameraset_cfg(),`
`57`	`57`	`max_relative_movement=np.deg2rad(5),`
`58`	`58`	`)`
	`59`	`+ env_rel.get_wrapper_attr("sim").open_gui()`
`59`	`60`
`60`	`61`	`for _ in range(10):`
`61`	`62`	`obs, info = env_rel.reset()`
Original file line number	Diff line number	Diff line change
`@@ -45,7 +45,7 @@ def create_sample_config_yaml(path: str):`
`45`	`45`	`camera_config=camera_cfg,`
`46`	`46`	`camera_type="realsense",`
`47`	`47`	`)`
`48`		`- cameras_sim = {"human_readable_name": SimCameraConfig(identifier="mjcf_name", type=0, on_screen_render=False)}`
	`48`	`+ cameras_sim = {"human_readable_name": SimCameraConfig(identifier="mjcf_name", type=0)}`
`49`	`49`	`sim = SimConfig(camera=SimCameraSetConfig(cameras=cameras_sim))`
`50`	`50`	`cfg = Config(hw=hw, sim=sim)`
`51`	`51`	`yml = to_yaml_str(cfg)`
Original file line number	Diff line number	Diff line change
`@@ -114,8 +114,8 @@ def default_fr3_sim_gripper_cfg():`
`114`	`114`	`def default_mujoco_cameraset_cfg():`
`115`	`115`
`116`	`116`	`cameras = {`
`117`		`- "wrist": SimCameraConfig(identifier="eye-in-hand_0", type=int(CameraType.fixed), on_screen_render=False),`
`118`		`- "default_free": SimCameraConfig(identifier="", type=int(CameraType.default_free), on_screen_render=True),`
	`117`	`+ "wrist": SimCameraConfig(identifier="eye-in-hand_0", type=int(CameraType.fixed)),`
	`118`	`+ "default_free": SimCameraConfig(identifier="", type=int(CameraType.default_free)),`
`119`	`119`	`}`
`120`	`120`	`return SimCameraSetConfig(cameras=cameras, resolution_width=1280, resolution_height=720, frame_rate=10)`
`121`	`121`