|
| 1 | +import atexit |
| 2 | +import multiprocessing as mp |
| 3 | +import uuid |
| 4 | +from logging import getLogger |
| 5 | +from multiprocessing.synchronize import Event as EventClass |
| 6 | +from os import PathLike |
| 7 | +from pathlib import Path |
| 8 | +from tempfile import NamedTemporaryFile |
| 9 | +from typing import Optional |
| 10 | + |
| 11 | +import mujoco as mj |
| 12 | +import mujoco.viewer |
| 13 | +from rcs._core.sim import GuiClient as _GuiClient |
| 14 | +from rcs._core.sim import Sim as _Sim |
| 15 | +from rcs.sim import egl_bootstrap |
| 16 | +from rcs.utils import SimpleFrameRate |
| 17 | + |
| 18 | +egl_bootstrap.bootstrap() |
| 19 | +logger = getLogger(__name__) |
| 20 | + |
| 21 | + |
| 22 | +# Target frames per second |
| 23 | +FPS = 60 |
| 24 | + |
| 25 | + |
| 26 | +def gui_loop(gui_uuid: str, close_event): |
| 27 | + frame_rate = SimpleFrameRate(FPS, "gui_loop") |
| 28 | + gui_client = _GuiClient(gui_uuid) |
| 29 | + model_bytes = gui_client.get_model_bytes() |
| 30 | + with NamedTemporaryFile(mode="wb") as f: |
| 31 | + f.write(model_bytes) |
| 32 | + model = mujoco.MjModel.from_binary_path(f.name) |
| 33 | + data = mujoco.MjData(model) |
| 34 | + gui_client.set_model_and_data(model._address, data._address) |
| 35 | + mujoco.mj_step(model, data) |
| 36 | + with mujoco.viewer.launch_passive(model, data) as viewer: |
| 37 | + while not close_event.is_set(): |
| 38 | + mujoco.mj_step(model, data) |
| 39 | + viewer.sync() |
| 40 | + gui_client.sync() |
| 41 | + frame_rate() |
| 42 | + |
| 43 | + |
| 44 | +class Sim(_Sim): |
| 45 | + def __init__(self, mjmdl: str | PathLike): |
| 46 | + mjmdl = Path(mjmdl) |
| 47 | + if mjmdl.suffix == ".xml": |
| 48 | + self.model = mj.MjModel.from_xml_path(str(mjmdl)) |
| 49 | + elif mjmdl.suffix == ".mjb": |
| 50 | + self.model = mj.MjModel.from_binary_path(str(mjmdl)) |
| 51 | + else: |
| 52 | + msg = f"Filetype {mjmdl.suffix} is unknown" |
| 53 | + logger.error(msg) |
| 54 | + self.data = mj.MjData(self.model) |
| 55 | + super().__init__(self.model._address, self.data._address) |
| 56 | + self._mp_context = mp.get_context("spawn") |
| 57 | + self._gui_uuid: Optional[str] = None |
| 58 | + self._gui_client: Optional[_GuiClient] = None |
| 59 | + self._gui_process: Optional[mp.context.SpawnProcess] = None |
| 60 | + self._stop_event: Optional[EventClass] = None |
| 61 | + |
| 62 | + def close_gui(self): |
| 63 | + if self._stop_event is not None: |
| 64 | + self._stop_event.set() |
| 65 | + if self._gui_process is not None: |
| 66 | + self._gui_process.join() |
| 67 | + self._stop_gui_server() |
| 68 | + |
| 69 | + def open_gui(self): |
| 70 | + if self._gui_uuid is None: |
| 71 | + self._gui_uuid = "rcs_" + str(uuid.uuid4()) |
| 72 | + self._start_gui_server(self._gui_uuid) |
| 73 | + if self._gui_client is None: |
| 74 | + ctx = mp.get_context("spawn") |
| 75 | + self._stop_event = ctx.Event() |
| 76 | + self._gui_process = ctx.Process( |
| 77 | + target=gui_loop, |
| 78 | + args=(self._gui_uuid, self._stop_event), |
| 79 | + ) |
| 80 | + self._gui_process.start() |
| 81 | + atexit.register(self.close_gui) |
0 commit comments