✨ better robot joint loading

src/runners/mjx.py

@@ -29,7 +29,7 @@ import numpy as np
 
 from src.core.env import BaseEnv
 from src.core.runner import BaseRunner, BaseRunnerConfig
-from src.runners.mujoco import MuJoCoRunner  # reuse _load_model_with_actuators
+from src.runners.mujoco import MuJoCoRunner  # reuse _load_model
 
 log = structlog.get_logger()
 
@@ -64,14 +64,8 @@ class MJXRunner(BaseRunner[MJXRunnerConfig]):
     # ── Initialization ───────────────────────────────────────────────
 
     def _sim_initialize(self, config: MJXRunnerConfig) -> None:
-        model_path = self.env.config.model_path
-        if model_path is None:
-            raise ValueError("model_path must be specified")
-
         # Step 1: Load CPU model (reuses URDF → MJCF → actuator injection)
-        self._mj_model = MuJoCoRunner._load_model_with_actuators(
-            str(model_path), self.env.config.actuators,
-        )
+        self._mj_model = MuJoCoRunner._load_model(self.env.robot)
         self._mj_model.opt.timestep = config.dt
         self._nq = self._mj_model.nq
         self._nv = self._mj_model.nv
@@ -207,10 +201,6 @@ class MJXRunner(BaseRunner[MJXRunnerConfig]):
         rv = self._batch_data.qvel[ids_np].astype(jnp.float32)
         return torch.from_dlpack(rq), torch.from_dlpack(rv)
 
-    def _sim_close(self) -> None:
-        if hasattr(self, "_offscreen_renderer") and self._offscreen_renderer is not None:
-            self._offscreen_renderer.close()
-
     # ── Rendering ────────────────────────────────────────────────────
 
     def render(self, env_idx: int = 0) -> np.ndarray: