VictorMylle/RL-Sim-Framework
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

✨ add mjx runner

Browse Source
This commit is contained in:
Victor Mylle
2026-03-09 21:18:19 +01:00
parent 15da0ef2fd
commit 26ccb1e902
8 changed files with 301 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
src/runners/mujoco.py
View File

@@ -198,12 +198,13 @@ class MuJoCoRunner(BaseRunner[MuJoCoRunnerConfig]):
def _sim_close(self) -> None:
if hasattr(self, "_offscreen_renderer") and self._offscreen_renderer is not None:
self._offscreen_renderer.close()
self._offscreen_renderer = None
self._data.clear()
def render(self, env_idx: int = 0) -> np.ndarray | None:
"""Offscreen render → RGB numpy array (H, W, 3)."""
if not hasattr(self, "_offscreen_renderer") or self._offscreen_renderer is None:
self._offscreen_renderer = mujoco.Renderer(self._model, height=480, width=640)
def render(self, env_idx: int = 0) -> np.ndarray:
"""Offscreen render of a single environment."""
if not hasattr(self, "_offscreen_renderer"):
self._offscreen_renderer = mujoco.Renderer(
self._model, width=640, height=480,
)
mujoco.mj_forward(self._model, self._data[env_idx])
self._offscreen_renderer.update_scene(self._data[env_idx])
return self._offscreen_renderer.render().copy()
return self._offscreen_renderer.render()