feat: SimBackend unified model property API

src/unilab/base/backend/base.py

@@ -1,4 +1,5 @@
 import abc
+from collections.abc import Sequence
 
 import numpy as np
 
@@ -26,6 +27,69 @@ def num_envs(self) -> int:
     def model(self):
         """底层物理模型"""
 
+    # ------------------------------------------------------------------ #
+    # Model properties                                                     #
+    # ------------------------------------------------------------------ #
+
+    @property
+    @abc.abstractmethod
+    def num_actuators(self) -> int:
+        """执行器数量"""
+
+    @property
+    @abc.abstractmethod
+    def num_dof_vel(self) -> int:
+        """关节速度自由度数量（不含浮动基座）"""
+
+    @abc.abstractmethod
+    def get_actuator_ctrl_range(self) -> np.ndarray:
+        """获取执行器控制范围
+
+        Returns:
+            (num_actuators, 2) 数组，列为 [low, high]
+        """
+
+    @abc.abstractmethod
+    def get_keyframe_qpos(self, name: str) -> np.ndarray:
+        """获取指定关键帧的完整 qpos（含浮动基座）
+
+        Args:
+            name: 关键帧名称（如 "stand"、"home"）
+
+        Returns:
+            (nq,) 数组
+        """
+
+    @abc.abstractmethod
+    def get_init_qvel(self) -> np.ndarray:
+        """获取零初始化的 qvel 向量，维度与 set_state 期望一致
+
+        Returns:
+            全零数组
+        """
+
+    @abc.abstractmethod
+    def get_body_ids(self, names: Sequence[str]) -> np.ndarray:
+        """将 body/link 名称解析为后端整数 ID
+
+        Args:
+            names: body/link 名称序列
+
+        Returns:
+            (len(names),) int32 数组
+
+        Raises:
+            ValueError: 若名称未找到
+        """
+
+    @abc.abstractmethod
+    def get_joint_range(self) -> np.ndarray | None:
+        """获取关节位置限制（不含浮动基座）
+
+        Returns:
+            (num_dof, 2) 数组，列为 [low, high]；若后端不支持则返回 None
+        """
+
     # ------------------------------------------------------------------ #
     # Simulation control                                                   #
     # ------------------------------------------------------------------ #

src/unilab/base/backend/motrix_backend.py

@@ -1,4 +1,5 @@
 import os
+from collections.abc import Sequence
 
 import numpy as np
 
@@ -106,6 +107,45 @@ def model(self):
     def data(self):
         return self._data
 
+    # ------------------------------------------------------------------ #
+    # Model properties                                                   #
+    # ------------------------------------------------------------------ #
+
+    @property
+    def num_actuators(self) -> int:
+        return int(self._model.num_actuators)
+
+    @property
+    def num_dof_vel(self) -> int:
+        # model.num_dof_vel includes 6 floating-base DOFs; exclude them
+        # to match get_dof_vel() which returns joint velocities only.
+        return int(self._model.num_dof_vel) - 6
+
+    def get_actuator_ctrl_range(self) -> np.ndarray:
+        arr: np.ndarray = np.array(self._model.actuator_ctrl_limits, dtype=self._np_dtype)
+        result: np.ndarray = arr.T.copy()
+        return result
+
+    def get_keyframe_qpos(self, name: str) -> np.ndarray:
+        if hasattr(self._model, "keyframes") and self._model.num_keyframes > 0:
+            return np.array(self._model.keyframes[0].dof_pos, dtype=self._np_dtype)
+        return np.array(self._model.compute_init_dof_pos(), dtype=self._np_dtype)
+
+    def get_init_qvel(self) -> np.ndarray:
+        return np.zeros((self._model.num_dof_vel,), dtype=self._np_dtype)
+
+    def get_body_ids(self, names: Sequence[str]) -> np.ndarray:
+        ids: list[int] = []
+        for name in names:
+            bid = self._model.get_link_index(name)
+            if bid is None or bid < 0:
+                raise ValueError(f"Body '{name}' not found in Motrix model")
+            ids.append(int(bid))
+        return np.array(ids, dtype=np.int32)
+
+    def get_joint_range(self) -> np.ndarray | None:
+        return None
+
     # ------------------------------------------------------------------ #
     # Simulation control                                                 #
     # ------------------------------------------------------------------ #

src/unilab/base/backend/mujoco_backend.py

@@ -1,4 +1,5 @@
 import os
+from collections.abc import Sequence
 from multiprocessing import cpu_count
 from typing import Optional, cast
 
@@ -187,6 +188,44 @@ def num_envs(self) -> int:
     def model(self):
         return self._model
 
+    # ------------------------------------------------------------------ #
+    # Model properties                                                   #
+    # ------------------------------------------------------------------ #
+
+    @property
+    def num_actuators(self) -> int:
+        return int(self._model.nu)
+
+    @property
+    def num_dof_vel(self) -> int:
+        return int(self._num_dof_vel)
+
+    def get_actuator_ctrl_range(self) -> np.ndarray:
+        return np.array(self._model.actuator_ctrlrange, dtype=self._np_dtype)
+
+    def get_keyframe_qpos(self, name: str) -> np.ndarray:
+        key_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_KEY, name)
+        if key_id < 0:
+            raise ValueError(f"Keyframe '{name}' not found in MuJoCo model")
+        return np.array(self._model.key_qpos[key_id].copy(), dtype=self._np_dtype)
+
+    def get_init_qvel(self) -> np.ndarray:
+        return np.zeros((self.nv,), dtype=self._np_dtype)
+
+    def get_body_ids(self, names: "Sequence[str]") -> np.ndarray:
+        ids: list[int] = []
+        for name in names:
+            bid = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_BODY, name)
+            if bid < 0:
+                raise ValueError(f"Body '{name}' not found in MuJoCo model")
+            ids.append(bid)
+        return np.array(ids, dtype=np.int32)
+
+    def get_joint_range(self) -> np.ndarray | None:
+        if self._root_qpos_dim > 0:
+            return np.array(self._model.jnt_range[1:], dtype=self._np_dtype)
+        return np.array(self._model.jnt_range, dtype=self._np_dtype)
+
     # ------------------------------------------------------------------ #
     # Simulation control                                                 #
     # ------------------------------------------------------------------ #

src/unilab/envs/locomotion/g1/base.py

@@ -3,7 +3,6 @@
 from dataclasses import dataclass, field
 
 import gymnasium as gym
-import mujoco
 import numpy as np
 
 from unilab.base.backend import SimBackend
@@ -61,36 +60,19 @@ def __init__(self, cfg: G1BaseCfg, backend: SimBackend, num_envs=1):
         self._init_buffers()
 
     def _init_action_space(self):
-        model = self._backend.model
-        if hasattr(model, "actuator_ctrlrange"):
-            low = model.actuator_ctrlrange[:, 0].copy()
-            high = model.actuator_ctrlrange[:, 1].copy()
-            nu = model.nu
-        else:
-            low = model.actuator_ctrl_limits[0, :]
-            high = model.actuator_ctrl_limits[1, :]
-            nu = model.num_actuators
-        self._action_space = gym.spaces.Box(low, high, (nu,), dtype=float)  # type: ignore[assignment]
+        ctrl_range = self._backend.get_actuator_ctrl_range()
+        nu = self._backend.num_actuators
+        self._action_space = gym.spaces.Box(ctrl_range[:, 0], ctrl_range[:, 1], (nu,), dtype=float)  # type: ignore[assignment]
 
     @property
     def action_space(self) -> gym.spaces.Box:
         return self._action_space  # type: ignore[no-any-return]
 
     def _init_buffers(self):
         self.default_angles = np.zeros((self._num_action,), dtype=np.float32)
-        model = self._backend.model
-        if hasattr(model, "key_qpos"):
-            key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, "stand")
-            if key_id >= 0:
-                self._init_qpos = np.array(model.key_qpos[key_id].copy(), dtype=np.float32)
-                self.default_angles = self._init_qpos[7:]
-            else:
-                raise ValueError("Keyframe 'stand' not found")
-            self._init_qvel = np.zeros((model.nv,), dtype=np.float32)
-        else:
-            self._init_qpos = model.compute_init_dof_pos()
-            self.default_angles = self._init_qpos[-self._num_action :]
-            self._init_qvel = np.zeros((model.num_dof_vel,), dtype=np.float32)
+        self._init_qpos = self._backend.get_keyframe_qpos("stand")
+        self.default_angles = self._init_qpos[-self._num_action :]
+        self._init_qvel = self._backend.get_init_qvel()
 
     def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray:
         state.info["last_actions"] = state.info.get("current_actions", actions.copy())

src/unilab/envs/locomotion/go1/base.py

@@ -3,7 +3,6 @@
 from dataclasses import dataclass, field
 
 import gymnasium as gym
-import mujoco
 import numpy as np
 
 from unilab.base.backend import SimBackend
@@ -80,16 +79,9 @@ def __init__(self, cfg: Go1BaseCfg, backend: SimBackend, num_envs=1):
         self._init_buffers()
 
     def _init_action_space(self):
-        model = self._backend.model
-        if hasattr(model, "actuator_ctrlrange"):
-            low = model.actuator_ctrlrange[:, 0].copy()
-            high = model.actuator_ctrlrange[:, 1].copy()
-            nu = model.nu
-        else:
-            low = model.actuator_ctrl_limits[0, :]
-            high = model.actuator_ctrl_limits[1, :]
-            nu = model.num_actuators
-        self._action_space = gym.spaces.Box(low, high, (nu,), dtype=float)  # type: ignore[assignment]
+        ctrl_range = self._backend.get_actuator_ctrl_range()
+        nu = self._backend.num_actuators
+        self._action_space = gym.spaces.Box(ctrl_range[:, 0], ctrl_range[:, 1], (nu,), dtype=float)  # type: ignore[assignment]
 
     @property
     def action_space(self) -> gym.spaces.Box:
@@ -98,26 +90,9 @@ def action_space(self) -> gym.spaces.Box:
     def _init_buffers(self):
         dtype = get_global_dtype()
         self.default_angles = np.zeros((self._num_action,), dtype=dtype)
-        model = self._backend.model
-        if hasattr(model, "key_qpos"):
-            # MuJoCo backend
-            key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, "home")
-            if key_id >= 0:
-                self._init_qpos = np.array(model.key_qpos[key_id].copy(), dtype=dtype)
-                self.default_angles = self._init_qpos[7:]
-            else:
-                raise ValueError("Keyframe 'home' not found in MuJoCo model")
-            self._init_qvel = np.zeros((model.nv,), dtype=dtype)
-        elif hasattr(model, "keyframes") and model.num_keyframes > 0:
-            # Motrix backend
-            kf = model.keyframes[0]  # Use first keyframe (should be "home")
-            self._init_qpos = np.array(kf.dof_pos, dtype=dtype)
-            self.default_angles = self._init_qpos[7:]
-            self._init_qvel = np.zeros((model.num_dof_vel,), dtype=dtype)
-        else:
-            raise ValueError(
-                "No keyframe found in model. Model must have either MuJoCo key_qpos or Motrix keyframes."
-            )
+        self._init_qpos = np.array(self._backend.get_keyframe_qpos("home"), dtype=dtype)
+        self.default_angles = self._init_qpos[7:]
+        self._init_qvel = self._backend.get_init_qvel().astype(dtype)
 
     def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray:
         state.info["last_actions"] = state.info.get("current_actions", np.zeros_like(actions))

src/unilab/envs/locomotion/go2/base.py

@@ -3,7 +3,6 @@
 from dataclasses import dataclass, field
 
 import gymnasium as gym
-import mujoco
 import numpy as np
 
 from unilab.base.backend import SimBackend
@@ -65,17 +64,9 @@ def __init__(self, cfg: Go2BaseCfg, backend: SimBackend, num_envs=1):
         self._init_buffers()
 
     def _init_action_space(self):
-        model = self._backend.model
-        if hasattr(model, "actuator_ctrlrange"):
-            low = model.actuator_ctrlrange[:, 0].copy()
-            high = model.actuator_ctrlrange[:, 1].copy()
-            nu = model.nu
-        else:
-            low = model.actuator_ctrl_limits[0, :]
-            high = model.actuator_ctrl_limits[1, :]
-            nu = model.num_actuators
-
-        self._action_space = gym.spaces.Box(low, high, (nu,), dtype=float)  # type: ignore[assignment]
+        ctrl_range = self._backend.get_actuator_ctrl_range()
+        nu = self._backend.num_actuators
+        self._action_space = gym.spaces.Box(ctrl_range[:, 0], ctrl_range[:, 1], (nu,), dtype=float)  # type: ignore[assignment]
 
     @property
     def action_space(self) -> gym.spaces.Box:
@@ -84,35 +75,9 @@ def action_space(self) -> gym.spaces.Box:
     def _init_buffers(self):
         dtype = get_global_dtype()
         self.default_angles = np.zeros((self._num_action,), dtype=np.float32)
-
-        model = self._backend.model
-        if hasattr(model, "key_qpos"):
-            # MuJoCo backend
-            key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, "home")
-            if key_id >= 0:
-                self._init_qpos = np.array(model.key_qpos[key_id].copy(), dtype=dtype)
-                self.default_angles = self._init_qpos[7:]
-            else:
-                raise ValueError("Keyframe 'home' not found in MuJoCo model")
-            self._init_qvel = np.zeros((model.nv,), dtype=dtype)
-        elif hasattr(model, "keyframes") and model.num_keyframes > 0:
-            # Motrix backend
-            kf = model.keyframes[0]  # Use first keyframe (should be "home")
-            self._init_qpos = np.array(kf.dof_pos, dtype=dtype)
-            self.default_angles = self._init_qpos[7:]
-            self._init_qvel = np.zeros((model.num_dof_vel,), dtype=dtype)
-        # if hasattr(model, "key_qpos"):  # MuJoCo
-        #     key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, "home")
-        #     if key_id >= 0:
-        #         self._init_qpos = np.array(model.key_qpos[key_id].copy(), dtype=np.float32)
-        #         self.default_angles = self._init_qpos[7:]
-        #     else:
-        #         raise ValueError("Keyframe 'home' not found")
-        #     self._init_qvel = np.zeros((model.nv,), dtype=np.float32)
-        # else:  # MotrixSim
-        #     self._init_qpos = model.compute_init_dof_pos()
-        #     self.default_angles = self._init_qpos[-self._num_action :]
-        #     self._init_qvel = np.zeros((model.num_dof_vel,), dtype=np.float32)
+        self._init_qpos = np.array(self._backend.get_keyframe_qpos("home"), dtype=dtype)
+        self.default_angles = self._init_qpos[7:]
+        self._init_qvel = self._backend.get_init_qvel().astype(dtype)
 
     def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray:
         state.info["last_actions"] = state.info.get("current_actions", actions.copy())

src/unilab/envs/motion_tracking/g1/tracking.py

@@ -313,24 +313,11 @@ def __init__(self, cfg: G1MotionTrackingCfg, num_envs=1, backend_type="mujoco"):
         self._log_action_scale_diagnostics()
 
         # Resolve body IDs for backend querying and motion-file indexing.
+        self.body_ids = self._backend.get_body_ids(cfg.body_names)
         if self._is_mujoco_backend():
-            self.body_ids = np.array(
-                [
-                    mujoco.mj_name2id(self._backend.model, mujoco.mjtObj.mjOBJ_BODY, name)
-                    for name in cfg.body_names
-                ],
-                dtype=np.int32,
-            )
             motion_body_ids = self.body_ids
         else:
-            backend_body_ids: list[int] = []
-            for name in cfg.body_names:
-                body_id = self._backend.model.get_link_index(name)
-                if body_id is None or body_id < 0:
-                    raise ValueError(f"Body '{name}' not found in motrix model")
-                backend_body_ids.append(int(body_id))
-            self.body_ids = np.array(backend_body_ids, dtype=np.int32)
-
+            # Motion data always uses MuJoCo body indexing.
             mj_model = mujoco.MjModel.from_xml_path(cfg.model_file)
             motion_body_ids = np.array(
                 [
@@ -376,9 +363,7 @@ def _get_body_pose_w(self) -> tuple[np.ndarray, np.ndarray]:
         )
 
     def _get_joint_range(self) -> np.ndarray | None:
-        if not self._is_mujoco_backend():
-            return None
-        return np.asarray(self._backend.model.jnt_range[1:])  # Skip free joint
+        return self._backend.get_joint_range()  # type: ignore[no-any-return]
 
     def _apply_adaptive_g1_action_scale(self) -> None:
         """Set per-joint action scale to match adaptive's normalization."""