src/runtime.py - master - koda/FiVI

src/runtime.py
276 строк10 KB
koda
Prepare production-ready grasp planner
18 июл 2026, 09:39
18 июл 2026, 09:39f540458
from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import mujoco  # type: ignore[import-untyped]
import numpy as np
import numpy.typing as npt


FloatArray = npt.NDArray[np.float64]
UInt8Array = npt.NDArray[np.uint8]
Int32Array = npt.NDArray[np.int32]

HAND_JOINT_COUNT: int = 20
DEFAULT_CAMERA_NAME: str = "hand_camera"
DEFAULT_CAMERA_WIDTH: int = 640
DEFAULT_CAMERA_HEIGHT: int = 480

HAND_JOINT_NAMES: tuple[str, ...] = (
    "rh_FFJ4", "rh_FFJ3", "rh_FFJ2", "rh_FFJ1",
    "rh_MFJ4", "rh_MFJ3", "rh_MFJ2", "rh_MFJ1",
    "rh_RFJ4", "rh_RFJ3", "rh_RFJ2", "rh_RFJ1",
    "rh_LFJ4", "rh_LFJ3", "rh_LFJ2", "rh_LFJ1",
    "rh_THJ5", "rh_THJ4", "rh_THJ2", "rh_THJ1",
)

HAND_ACTUATOR_NAMES: tuple[str, ...] = (
    "rh_A_FFJ4", "rh_A_FFJ3", "rh_A_FFJ0",
    "rh_A_MFJ4", "rh_A_MFJ3", "rh_A_MFJ0",
    "rh_A_RFJ4", "rh_A_RFJ3", "rh_A_RFJ0",
    "rh_A_LFJ4", "rh_A_LFJ3", "rh_A_LFJ0",
    "rh_A_THJ5", "rh_A_THJ4", "rh_A_THJ2", "rh_A_THJ1",
)

@dataclass(frozen=True, slots=True)
class RuntimeConfig:
    """Immutable MuJoCo model and camera configuration."""

    xml_path: Path
    camera_name: str = DEFAULT_CAMERA_NAME
    camera_width: int = DEFAULT_CAMERA_WIDTH
    camera_height: int = DEFAULT_CAMERA_HEIGHT
    xml_string: str | None = None

    def __post_init__(self) -> None:
        path = Path(self.xml_path).expanduser().resolve()
        if not path.is_file():
            raise ValueError(f"MuJoCo XML model does not exist: {path}")
        if path.suffix.lower() != ".xml":
            raise ValueError(f"MuJoCo model must be an XML file: {path}")
        if not self.camera_name:
            raise ValueError("camera_name must not be empty")
        if self.camera_width <= 0 or self.camera_height <= 0:
            raise ValueError("camera dimensions must be positive")
        if self.xml_string is not None and not self.xml_string.strip():
            raise ValueError("xml_string must be non-empty when provided")
        object.__setattr__(self, "xml_path", path)


class Runtime:
    """Own an ``MjModel``, ``MjData``, and off-screen RGB-D renderer."""

    def __init__(self, config: RuntimeConfig) -> None:
        """Load the configured XML model and allocate runtime state."""

        self._config = config
        try:
            self._model = (
                mujoco.MjModel.from_xml_path(str(config.xml_path))
                if config.xml_string is None
                else mujoco.MjModel.from_xml_string(config.xml_string)
            )
        except Exception as error:
            raise ValueError(
                f"failed to load MuJoCo XML model {config.xml_path}: {error}"
            ) from error

        camera_id = mujoco.mj_name2id(
            self._model, mujoco.mjtObj.mjOBJ_CAMERA, config.camera_name
        )
        if camera_id < 0:
            raise ValueError(f"camera not found in MuJoCo model: {config.camera_name}")

        self._hand_qpos_indices, self._hand_dof_indices = self._resolve_hand_joints()
        self._hand_actuator_ids = self._resolve_hand_actuators()
        self._data = mujoco.MjData(self._model)
        try:
            self._renderer = mujoco.Renderer(
                self._model,
                height=config.camera_height,
                width=config.camera_width,
            )
        except Exception as error:
            raise ValueError(f"failed to initialize MuJoCo renderer: {error}") from error
        self._closed = False
        mujoco.mj_forward(self._model, self._data)

    def _ensure_open(self) -> None:
        if self._closed:
            raise RuntimeError("runtime is closed")

    def _required_id(self, object_type: mujoco.mjtObj, name: str) -> int:
        object_id = mujoco.mj_name2id(self._model, object_type, name)
        if object_id < 0:
            raise ValueError(f"MuJoCo object not found: {name}")
        return int(object_id)

    def _resolve_hand_joints(
        self,
    ) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]:
        qpos_indices = np.empty(HAND_JOINT_COUNT, dtype=np.intp)
        dof_indices = np.empty(HAND_JOINT_COUNT, dtype=np.intp)
        for index, joint_name in enumerate(HAND_JOINT_NAMES):
            joint_id = self._required_id(mujoco.mjtObj.mjOBJ_JOINT, joint_name)
            qpos_indices[index] = int(self._model.jnt_qposadr[joint_id])
            dof_indices[index] = int(self._model.jnt_dofadr[joint_id])
        qpos_indices.flags.writeable = False
        dof_indices.flags.writeable = False
        return qpos_indices, dof_indices

    def _resolve_hand_actuators(self) -> npt.NDArray[np.intp]:
        actuator_ids = np.array(
            tuple(
                self._required_id(mujoco.mjtObj.mjOBJ_ACTUATOR, name)
                for name in HAND_ACTUATOR_NAMES
            ),
            dtype=np.intp,
        )
        actuator_ids.flags.writeable = False
        return actuator_ids

    @property
    def model(self) -> mujoco.MjModel:
        """Return the loaded read-mostly MuJoCo model."""

        return self._model

    @property
    def data(self) -> mujoco.MjData:
        """Return the mutable MuJoCo simulation data."""

        return self._data

    @property
    def hand_qpos_indices(self) -> npt.NDArray[np.intp]:
        """Return immutable qpos addresses for the twenty hand joints."""

        return self._hand_qpos_indices

    @property
    def hand_dof_indices(self) -> npt.NDArray[np.intp]:
        """Return immutable qvel addresses for the twenty hand joints."""

        return self._hand_dof_indices

    @property
    def hand_actuator_ids(self) -> npt.NDArray[np.intp]:
        """Return immutable actuator IDs in reduced sixteen-DOF order."""

        return self._hand_actuator_ids

    def get_joint_positions(self) -> FloatArray:
        """Return a detached ``float64`` copy of ``qpos`` with shape ``(20,)``."""

        positions = np.array(
            self._data.qpos[self._hand_qpos_indices], dtype=np.float64, copy=True
        )
        if positions.shape != (HAND_JOINT_COUNT,):
            raise ValueError(f"unexpected qpos shape: {positions.shape}")
        return positions

    def set_joint_positions(self, positions: npt.ArrayLike) -> None:
        """Set the full physical ``qpos`` state and recompute derived data.

        This is a state setter, not an actuator command. Position actuator
        targets remain in ``data.ctrl`` and are not modified.
        """

        values = np.asarray(positions, dtype=np.float64)
        if values.shape != (HAND_JOINT_COUNT,):
            raise ValueError(
                f"positions must have shape ({HAND_JOINT_COUNT},), got {values.shape}"
            )
        if not bool(np.all(np.isfinite(values))):
            raise ValueError("positions must contain only finite values")
        self._data.qpos[self._hand_qpos_indices] = values
        mujoco.mj_forward(self._model, self._data)

    def get_camera_image(self) -> tuple[UInt8Array, FloatArray]:
        """Render the configured camera and return detached RGB and depth arrays."""

        self._ensure_open()
        self._renderer.update_scene(self._data, camera=self._config.camera_name)
        self._renderer.disable_depth_rendering()
        rgb = np.array(self._renderer.render(), dtype=np.uint8, copy=True)

        self._renderer.enable_depth_rendering()
        try:
            depth = np.array(self._renderer.render(), dtype=np.float64, copy=True)
        finally:
            self._renderer.disable_depth_rendering()

        expected_rgb = (
            self._config.camera_height,
            self._config.camera_width,
            3,
        )
        expected_depth = (
            self._config.camera_height,
            self._config.camera_width,
        )
        if rgb.shape != expected_rgb:
            raise ValueError(f"unexpected RGB image shape: {rgb.shape}")
        if depth.shape != expected_depth:
            raise ValueError(f"unexpected depth image shape: {depth.shape}")
        return rgb, depth

    def get_camera_segmentation(self) -> Int32Array:
        """Render detached MuJoCo object IDs from the configured camera.

        Segmentation IDs are used only to isolate camera pixels. Object pose,
        dimensions, material, and dynamics are not read by this method.
        """

        self._ensure_open()
        self._renderer.update_scene(self._data, camera=self._config.camera_name)
        self._renderer.enable_segmentation_rendering()
        try:
            segmentation = np.array(self._renderer.render(), dtype=np.int32, copy=True)
        finally:
            self._renderer.disable_segmentation_rendering()
        expected = (self._config.camera_height, self._config.camera_width, 2)
        if segmentation.shape != expected:
            raise ValueError(
                f"unexpected segmentation image shape: {segmentation.shape}"
            )
        return segmentation

    def close(self) -> None:
        """Release the off-screen renderer resources."""

        if self._closed:
            return
        try:
            self._renderer.close()
        finally:
            self._closed = True

    def __enter__(self) -> Runtime:
        """Return this runtime as a context manager."""

        return self

    def __exit__(
        self,
        exception_type: type[BaseException] | None,
        exception: BaseException | None,
        traceback: object | None,
    ) -> None:
        """Release renderer resources when leaving a context manager."""

        del exception_type, exception, traceback
        self.close()


__all__ = [
    "DEFAULT_CAMERA_HEIGHT",
    "DEFAULT_CAMERA_NAME",
    "DEFAULT_CAMERA_WIDTH",
    "FloatArray",
    "HAND_JOINT_COUNT",
    "Runtime",
    "RuntimeConfig",
    "UInt8Array",
]