From c07aed03314e9ee6a8413c7b626579752139cef9 Mon Sep 17 00:00:00 2001 From: Nabla7 Date: Thu, 25 Jun 2026 02:07:02 +0200 Subject: [PATCH 01/73] Add runtime scene package loading --- data/.lfs/scene_packages.tar.gz | 3 + dimos/control/coordinator.py | 27 ++ .../g1_groot_wbc_task/g1_groot_wbc_task.py | 62 +++- .../test_g1_groot_wbc_task.py | 32 ++ dimos/control/test_control.py | 38 ++ dimos/core/global_config.py | 1 + dimos/hardware/whole_body/registry.py | 23 +- dimos/robot/unitree/g1/assets/g1_29dof.xml | 297 +++++++++++++++ .../blueprints/basic/unitree_g1_groot_wbc.py | 64 +++- dimos/simulation/engines/mujoco_engine.py | 174 ++++++++- dimos/simulation/engines/mujoco_sim_module.py | 156 +++++++- .../engines/test_mujoco_sim_module.py | 221 ++++++++++- dimos/simulation/mujoco/entity_scene.py | 347 ++++++++++++++++++ dimos/simulation/scene_assets/spec.py | 276 ++++++++++++++ dimos/simulation/scenes/catalog.py | 134 +++++++ dimos/simulation/utils/xml_parser.py | 6 +- stubs/mujoco/__init__.pyi | 25 ++ 17 files changed, 1856 insertions(+), 30 deletions(-) create mode 100644 data/.lfs/scene_packages.tar.gz create mode 100644 dimos/robot/unitree/g1/assets/g1_29dof.xml create mode 100644 dimos/simulation/mujoco/entity_scene.py create mode 100644 dimos/simulation/scene_assets/spec.py create mode 100644 dimos/simulation/scenes/catalog.py diff --git a/data/.lfs/scene_packages.tar.gz b/data/.lfs/scene_packages.tar.gz new file mode 100644 index 0000000000..32e9137c05 --- /dev/null +++ b/data/.lfs/scene_packages.tar.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bb91986a9dc14759a7f767c00b54416e9d6eb09f1c053c24a47a3ea6fe20cd1b +size 707835985 diff --git a/dimos/control/coordinator.py b/dimos/control/coordinator.py index 54022f5892..fb2bfbdc6c 100644 --- a/dimos/control/coordinator.py +++ b/dimos/control/coordinator.py @@ -26,6 +26,7 @@ """ from dataclasses import dataclass, field +import inspect import threading import time from typing import TYPE_CHECKING, Any @@ -551,6 +552,32 @@ def set_dry_run(self, enabled: bool) -> None: except Exception: logger.exception(f"set_dry_run() raised on task {task.name!r}") + @rpc + def reset_runtime_state(self, reactivate: bool | None = None) -> dict[str, bool]: + """Reset transient state on tasks that expose ``reset_runtime_state``. + + This is meant for simulation/runtime discontinuities such as MuJoCo + respawn, where task histories and latched commands must be cleared + without tearing down the coordinator. + """ + results: dict[str, bool] = {} + with self._task_lock: + for task in self._tasks.values(): + handler = getattr(task, "reset_runtime_state", None) + if not callable(handler): + results[task.name] = False + continue + try: + params = inspect.signature(handler).parameters + if "reactivate" in params: + results[task.name] = bool(handler(reactivate=reactivate)) + else: + results[task.name] = bool(handler()) + except Exception: + logger.exception(f"reset_runtime_state() raised on task {task.name!r}") + results[task.name] = False + return results + @rpc def task_invoke( self, task_name: TaskName, method: str, kwargs: dict[str, Any] | None = None diff --git a/dimos/control/tasks/g1_groot_wbc_task/g1_groot_wbc_task.py b/dimos/control/tasks/g1_groot_wbc_task/g1_groot_wbc_task.py index 98972195db..9ca032dc40 100644 --- a/dimos/control/tasks/g1_groot_wbc_task/g1_groot_wbc_task.py +++ b/dimos/control/tasks/g1_groot_wbc_task/g1_groot_wbc_task.py @@ -171,6 +171,26 @@ _NUM_MOTORS = 29 +def _preferred_onnx_providers() -> list[str]: + available = ort.get_available_providers() + providers: list[str] = [] + + if "CUDAExecutionProvider" in available: + preload_dlls = getattr(ort, "preload_dlls", None) + if preload_dlls is not None: + try: + preload_dlls(cuda=True, cudnn=True, msvc=False) + except Exception as exc: + logger.warning( + "Failed to preload ONNXRuntime CUDA/cuDNN libraries", + error=repr(exc), + ) + providers.append("CUDAExecutionProvider") + + providers.append("CPUExecutionProvider") + return providers + + @dataclass class G1GrootWBCTaskConfig: """Configuration for the GR00T WBC task. @@ -289,7 +309,7 @@ def __init__( self._joint_names_set = frozenset(config.joint_names) self._all_joint_names = list(config.all_joint_names) - providers = ort.get_available_providers() + providers = _preferred_onnx_providers() self._balance_session = ort.InferenceSession(str(config.balance_onnx), providers=providers) self._walk_session = ort.InferenceSession(str(config.walk_onnx), providers=providers) self._balance_input = self._balance_session.get_inputs()[0].name @@ -299,7 +319,9 @@ def __init__( task=name, balance=str(config.balance_onnx), walk=str(config.walk_onnx), - providers=providers, + requested_providers=providers, + balance_providers=self._balance_session.get_providers(), + walk_providers=self._walk_session.get_providers(), ) self._default_29 = np.asarray(config.default_positions_29, dtype=np.float32) @@ -663,6 +685,42 @@ def disarm(self) -> bool: logger.info("G1GrootWBCTask disarmed (holding current pose)", task=self._name) return True + def reset_runtime_state(self, reactivate: bool | None = None) -> bool: + """Clear runtime policy state after a simulation discontinuity. + + ``reactivate=None`` preserves whether the task was armed/arming before + reset. Passing ``True`` forces a clean immediate re-arm on the next + coordinator tick, which is useful after MuJoCo respawn. + """ + was_armed = self._armed or self._arming or self._arm_pending + should_reactivate = was_armed if reactivate is None else bool(reactivate) + + self._armed = False + self._arming = False + self._arm_pending = False + self._ramp_start = None + self._arming_start_t = 0.0 + self._last_targets = None + self._state_seen = False + self._cached_q_29[:] = self._default_29 + self._cached_dq_29[:] = 0.0 + self._cached_q_15[:] = self._default_15 + self._reset_policy_state() + with self._cmd_lock: + self._cmd[:] = 0.0 + self._last_cmd_time = 0.0 + + if self._active and should_reactivate: + self._arming_duration = 0.0 + self._arm_pending = True + + logger.info( + "G1GrootWBCTask runtime state reset", + task=self._name, + reactivate=should_reactivate, + ) + return True + def set_dry_run(self, enabled: bool) -> None: """Enable/disable dry-run. diff --git a/dimos/control/tasks/g1_groot_wbc_task/test_g1_groot_wbc_task.py b/dimos/control/tasks/g1_groot_wbc_task/test_g1_groot_wbc_task.py index 2e11da90a4..0359a7a41a 100644 --- a/dimos/control/tasks/g1_groot_wbc_task/test_g1_groot_wbc_task.py +++ b/dimos/control/tasks/g1_groot_wbc_task/test_g1_groot_wbc_task.py @@ -48,11 +48,13 @@ def __init__( label: str, action: np.ndarray, call_log: list[str], + providers: Any = None, ) -> None: self.model_path = model_path self._label = label self._action = action self._call_log = call_log + self._providers = list(providers or []) fake_input = MagicMock() fake_input.name = "obs" self._inputs = [fake_input] @@ -60,6 +62,9 @@ def __init__( def get_inputs(self) -> list[Any]: return self._inputs + def get_providers(self) -> list[str]: + return self._providers + def run(self, _outputs: Any, _feed: dict[str, np.ndarray]) -> list[np.ndarray]: self._call_log.append(self._label) return [self._action.reshape(1, -1)] @@ -77,6 +82,7 @@ def _factory(path: str, providers: Any = None) -> _StubSession: label=label, action=np.full(15, 0.1, dtype=np.float32), call_log=call_log, + providers=providers, ) monkeypatch.setattr(g1_groot_wbc_task.ort, "InferenceSession", _factory) @@ -278,6 +284,32 @@ def test_dry_run_suppresses_output_but_keeps_policy_hot( assert np.any(task._obs_buf != 0.0) +def test_reset_runtime_state_clears_policy_state_and_rearms( + task: G1GrootWBCTask, joints_29: list[str] +) -> None: + task.start() + task.set_velocity_command(0.5, 0.0, 0.0, t_now=100.0) + for _ in range(10): + task.compute(_state_at(100.0, joints_29)) + + assert task.state_snapshot()["armed"] + assert np.any(task._obs_buf != 0.0) + assert np.any(task._cmd != 0.0) + + assert task.reset_runtime_state(reactivate=True) + + snapshot = task.state_snapshot() + assert not snapshot["armed"] + assert snapshot["arm_pending"] + np.testing.assert_array_equal(task._obs_buf, np.zeros_like(task._obs_buf)) + np.testing.assert_array_equal(task._last_action, np.zeros_like(task._last_action)) + np.testing.assert_array_equal(task._cmd, np.zeros_like(task._cmd)) + assert task._last_cmd_time == 0.0 + + task.compute(_state_at(101.0, joints_29)) + assert task.state_snapshot()["armed"] + + def test_projected_gravity_matches_reference_quaternion_order() -> None: np.testing.assert_allclose( G1GrootWBCTask._projected_gravity((1.0, 0.0, 0.0, 0.0)), diff --git a/dimos/control/test_control.py b/dimos/control/test_control.py index ae6bc1e9de..c34988c0e9 100644 --- a/dimos/control/test_control.py +++ b/dimos/control/test_control.py @@ -26,6 +26,7 @@ from dimos.control.coordinator import ControlCoordinator from dimos.control.hardware_interface import ConnectedHardware from dimos.control.task import ( + BaseControlTask, ControlMode, CoordinatorState, JointCommandOutput, @@ -190,6 +191,43 @@ def test_write_command(self, connected_hardware, mock_adapter): class TestControlCoordinatorLifecycle: + def test_reset_runtime_state_calls_task_hooks(self): + class ResettableTask(BaseControlTask): + def __init__(self) -> None: + self.reset_reactivate_args: list[bool | None] = [] + + @property + def name(self) -> str: + return "resettable" + + def claim(self) -> ResourceClaim: + return ResourceClaim(joints=frozenset()) + + def is_active(self) -> bool: + return True + + def compute(self, state: CoordinatorState) -> JointCommandOutput | None: + _ = state + return None + + def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: + _ = by_task, joints + + def reset_runtime_state(self, reactivate: bool | None = None) -> bool: + self.reset_reactivate_args.append(reactivate) + return True + + coordinator = ControlCoordinator(publish_joint_state=False) + task = ResettableTask() + + try: + assert coordinator.add_task(task) + + assert coordinator.reset_runtime_state(reactivate=True) == {"resettable": True} + assert task.reset_reactivate_args == [True] + finally: + coordinator.stop() + def test_start_stop_calls_adapter_activate_and_deactivate(self): from dimos.hardware.manipulators.mock.adapter import MockAdapter from dimos.hardware.manipulators.registry import adapter_registry diff --git a/dimos/core/global_config.py b/dimos/core/global_config.py index c75ea6fd18..9d287c0f59 100644 --- a/dimos/core/global_config.py +++ b/dimos/core/global_config.py @@ -58,6 +58,7 @@ class GlobalConfig(BaseSettings): mujoco_global_map_from_pointcloud: str | None = None mujoco_start_pos: str = "-1.0, 1.0" mujoco_steps_per_frame: int = 7 + scene: str | None = None robot_model: str | None = None robot_id: str | None = None robot_width: float = 0.3 diff --git a/dimos/hardware/whole_body/registry.py b/dimos/hardware/whole_body/registry.py index ec841fe182..4480de6ebe 100644 --- a/dimos/hardware/whole_body/registry.py +++ b/dimos/hardware/whole_body/registry.py @@ -71,14 +71,33 @@ def available(self) -> list[str]: return sorted(self._adapters.keys()) def discover(self) -> None: - """Discover and register whole-body hardware adapters. + """Discover and register whole-body adapters. Walks the hardware whole-body package recursively looking for - ``adapter.py`` modules that provide a ``register(registry)`` function. + ``adapter.py`` modules, then scans simulation whole-body modules. Any + discovered module can provide a ``register(registry)`` function. """ import dimos.hardware.whole_body as hw_pkg + import dimos.simulation.adapters.whole_body as sim_pkg self._discover_in("dimos.hardware.whole_body", hw_pkg.__path__[0], max_depth=2) + self._discover_simulation_adapters( + "dimos.simulation.adapters.whole_body", + sim_pkg.__path__[0], + ) + + def _discover_simulation_adapters(self, pkg_path: str, dir_path: str) -> None: + for entry in sorted(os.listdir(dir_path)): + if entry.startswith(("_", ".")) or not entry.endswith(".py"): + continue + module_name = entry.removesuffix(".py") + try: + mod = importlib.import_module(f"{pkg_path}.{module_name}") + except ImportError as e: + logger.warning(f"Skipping whole-body simulation adapter {module_name}: {e}") + continue + if hasattr(mod, "register"): + mod.register(self) def _discover_in(self, pkg_path: str, dir_path: str, *, max_depth: int) -> None: for entry in sorted(os.listdir(dir_path)): diff --git a/dimos/robot/unitree/g1/assets/g1_29dof.xml b/dimos/robot/unitree/g1/assets/g1_29dof.xml new file mode 100644 index 0000000000..b1dd123152 --- /dev/null +++ b/dimos/robot/unitree/g1/assets/g1_29dof.xml @@ -0,0 +1,297 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py index 3428ac91a1..1709930830 100644 --- a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py @@ -34,6 +34,8 @@ Usage: dimos run unitree-g1-groot-wbc # real hardware dimos --simulation mujoco run unitree-g1-groot-wbc # sim + dimos --simulation mujoco --scene data/scene_packages/dimos_office/scene.meta.json \\ + run unitree-g1-groot-wbc # sim in a cooked scene package Overrides (replace the old env-var dance): dimos run unitree-g1-groot-wbc \\ @@ -70,8 +72,11 @@ # whole CLI on a multi-GB download every time the module is imported. _GROOT_MODEL_DIR = LfsPath("groot") _MJCF_PATH = LfsPath("mujoco_sim/g1_gear_wbc.xml") +_ROBOT_ONLY_MJCF_PATH = Path(__file__).resolve().parents[2] / "assets" / "g1_29dof.xml" +_ROBOT_MESHDIR = LfsPath("g1_urdf/meshes") _adapter_address: str | Path +_cmd_vel_topic = "/cmd_vel" if global_config.simulation else "/g1/cmd_vel" if global_config.simulation and global_config.simulation != "mujoco": raise ValueError("unitree-g1-groot-wbc only supports --simulation mujoco") @@ -94,35 +99,68 @@ imu_gyro_names=( "imu-pelvis-angular-velocity", "imu-torso-angular-velocity", + "imu-angular-velocity", "gyro_pelvis", "imu_gyro", ), imu_accel_names=( "imu-pelvis-linear-acceleration", "imu-torso-linear-acceleration", + "imu-linear-acceleration", "accelerometer_pelvis", "imu_accel", ), require_imu=True, ) + def _legacy_mujoco_backend() -> Any: + return MujocoSimModule.blueprint( + address=_MJCF_PATH, + headless=True, + dof=29, + enable_color=False, + enable_depth=False, + enable_pointcloud=False, + inject_legacy_assets=True, + robot_sim_spec=_g1_sim_spec, + ) + + def _scene_mujoco_backend() -> tuple[Any, str | Path]: + if global_config.scene is None: + return _legacy_mujoco_backend(), _MJCF_PATH + + from dimos.simulation.scenes.catalog import resolve_scene_package + + package = resolve_scene_package(global_config.scene) + if package is None: + return _legacy_mujoco_backend(), _MJCF_PATH + if package.mujoco_scene_path is None: + raise ValueError(f"scene package has no MuJoCo scene artifact: {package.metadata_path}") + + return ( + MujocoSimModule.blueprint( + scene_xml=package.mujoco_scene_path, + robot_mjcf=_ROBOT_ONLY_MJCF_PATH, + robot_meshdir=_ROBOT_MESHDIR, + robot_id="", + scene_entities=package.entities, + headless=True, + dof=29, + enable_color=False, + enable_depth=False, + enable_pointcloud=False, + robot_sim_spec=_g1_sim_spec, + ), + _ROBOT_ONLY_MJCF_PATH, + ) + # Sim backend: MuJoCo engine via SHM. - _backend = MujocoSimModule.blueprint( - address=_MJCF_PATH, - headless=True, - dof=29, - enable_color=False, - enable_depth=False, - enable_pointcloud=False, - inject_legacy_assets=True, - robot_sim_spec=_g1_sim_spec, - ) + _backend, _adapter_address = _scene_mujoco_backend() # MujocoSimModule's ``odom`` Out is the sole producer of ``/odom`` # now - the coordinator no longer polls the whole-body adapter for # base pose (read_odom was dropped from the Protocol). autoconnect # maps ``(odom, PoseStamped)`` to ``/odom`` by default; no override. _adapter_type = "sim_mujoco_g1" - _adapter_address = _MJCF_PATH _tick_rate = 50.0 _auto_arm = True _auto_dry_run = False @@ -246,8 +284,8 @@ def _viewer() -> Any: ).transports( { ("joint_command", JointState): LCMTransport("/g1/joint_command", JointState), - ("twist_command", Twist): LCMTransport("/g1/cmd_vel", Twist), - ("tele_cmd_vel", Twist): LCMTransport("/g1/cmd_vel", Twist), + ("twist_command", Twist): LCMTransport(_cmd_vel_topic, Twist), + ("tele_cmd_vel", Twist): LCMTransport(_cmd_vel_topic, Twist), # Real-hw only: the transport_lcm adapter speaks to # G1WholeBodyConnection over these topics. autoconnect already # matches by (name, type) so sim doesn't need them -- they're diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index d9889027f4..0454a6bda1 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -18,6 +18,7 @@ from collections.abc import Callable from dataclasses import dataclass +import math from pathlib import Path import threading import time @@ -47,6 +48,7 @@ StepHook = Callable[["MujocoEngine"], None] _MJJNT_FREE = int(mujoco.mjtJoint.mjJNT_FREE) # type: ignore[attr-defined] +_RESET_WAIT_TIMEOUT_S = 5.0 @dataclass @@ -94,14 +96,27 @@ def __init__( on_before_step: StepHook | None = None, on_after_step: StepHook | None = None, assets: dict[str, bytes] | None = None, + model: mujoco.MjModel | None = None, robot_sim_spec: RobotSimSpec | None = None, + spawn_xy: tuple[float, float] | None = None, + spawn_z: float | None = None, + spawn_yaw: float | None = None, + reset_joint_positions: list[float] | None = None, ) -> None: super().__init__(config_path=config_path, headless=headless) self._on_before_step: StepHook | None = on_before_step self._on_after_step: StepHook | None = on_after_step + self._spawn_xy = spawn_xy + self._spawn_z = spawn_z + self._spawn_yaw = spawn_yaw + self._reset_joint_positions = reset_joint_positions xml_path = self._resolve_xml_path(config_path) - if assets is not None: + if model is not None: + if assets is not None: + raise ValueError("MujocoEngine cannot use injected assets with a precompiled model") + self._model = model + elif assets is not None: # MJCFs that reference meshes by bare filename (e.g. menagerie # G1) need the mesh bytes injected by name; from_xml_path can't # find them on disk. @@ -113,7 +128,8 @@ def __init__( self._xml_path = xml_path self._data = mujoco.MjData(self._model) - self._joint_mappings = build_joint_mappings(self._xml_path, self._model) + mapping_xml_path = None if model is not None else self._xml_path + self._joint_mappings = build_joint_mappings(mapping_xml_path, self._model) self._robot_binding: RobotSimBinding | None = None if robot_sim_spec is not None: self._robot_binding = resolve_robot_sim_binding( @@ -131,6 +147,8 @@ def __init__( self._connected = False self._lock = threading.Lock() + self._reset_requested = threading.Event() + self._reset_done_event: threading.Event | None = None self._stop_event = threading.Event() self._sim_thread: threading.Thread | None = None @@ -142,6 +160,8 @@ def __init__( self._joint_velocity_targets = [0.0] * self._num_joints self._joint_effort_targets = [0.0] * self._num_joints self._command_mode = "position" + self._apply_spawn_pose_unlocked() + self._apply_reset_joint_positions_unlocked() for i, mapping in enumerate(self._joint_mappings): current_pos = self._current_position(mapping) self._joint_position_targets[i] = current_pos @@ -318,6 +338,63 @@ def _close_cam_renderers(cam_renderers: dict[str, _CameraRendererState]) -> None state.rgb_renderer.close() state.depth_renderer.close() + def _reset_unlocked(self) -> None: + if self._model.nkey > 0: + mujoco.mj_resetDataKeyframe(self._model, self._data, 0) + else: + mujoco.mj_resetData(self._model, self._data) + self._apply_spawn_pose_unlocked() + self._apply_reset_joint_positions_unlocked() + for i, mapping in enumerate(self._joint_mappings): + self._joint_position_targets[i] = self._current_position(mapping) + self._command_mode = "position" + + root_pose = self.get_root_pose_unlocked() + if root_pose is not None: + position, _ = root_pose + logger.info( + "MuJoCo reset applied", + x=float(position[0]), + y=float(position[1]), + z=float(position[2]), + ) + + def _apply_spawn_pose_unlocked(self) -> None: + qpos_adr = self._root_qpos_adr + if qpos_adr is None: + mujoco.mj_forward(self._model, self._data) + return + + qpos = self._data.qpos + if self._spawn_xy is not None: + qpos[qpos_adr] = float(self._spawn_xy[0]) + qpos[qpos_adr + 1] = float(self._spawn_xy[1]) + if self._spawn_z is not None: + qpos[qpos_adr + 2] = float(self._spawn_z) + if self._spawn_yaw is not None: + qpos[qpos_adr + 3 : qpos_adr + 7] = [ + math.cos(float(self._spawn_yaw) * 0.5), + 0.0, + 0.0, + math.sin(float(self._spawn_yaw) * 0.5), + ] + + qvel_adr = self._root_qvel_adr + if qvel_adr is not None: + self._data.qvel[qvel_adr : qvel_adr + 6] = 0.0 + mujoco.mj_forward(self._model, self._data) + + def _apply_reset_joint_positions_unlocked(self) -> None: + if self._reset_joint_positions is None: + return + for index, position in enumerate(self._reset_joint_positions[: self._num_joints]): + mapping = self._joint_mappings[index] + if mapping.qpos_adr is not None: + self._data.qpos[mapping.qpos_adr] = float(position) + if mapping.dof_adr is not None: + self._data.qvel[mapping.dof_adr] = 0.0 + mujoco.mj_forward(self._model, self._data) + def _sim_loop(self) -> None: logger.info("sim loop started", cls=self.__class__.__name__) dt = 1.0 / self._control_frequency @@ -327,6 +404,15 @@ def _sim_loop(self) -> None: def _step_once(sync_viewer: bool) -> None: loop_start = time.time() + reset_done_event = None + if self._reset_requested.is_set(): + with self._lock: + self._reset_requested.clear() + self._reset_unlocked() + reset_done_event = self._reset_done_event + self._reset_done_event = None + if reset_done_event is not None: + reset_done_event.set() if self._on_before_step is not None: try: self._on_before_step(self) @@ -486,6 +572,90 @@ def hold_current_position(self) -> None: for i, mapping in enumerate(self._joint_mappings): self._joint_position_targets[i] = self._current_position(mapping) + def reset(self) -> None: + with self._lock: + self._reset_unlocked() + + def request_reset( + self, + *, + wait: bool = False, + timeout: float = _RESET_WAIT_TIMEOUT_S, + ) -> bool: + done_event = threading.Event() if wait else None + with self._lock: + self._reset_done_event = done_event + self._reset_requested.set() + if done_event is None: + return True + return done_event.wait(timeout) + + def request_reset_to( + self, + *, + spawn_xy: tuple[float, float], + spawn_z: float | None = None, + spawn_yaw: float | None = None, + wait: bool = False, + timeout: float = _RESET_WAIT_TIMEOUT_S, + ) -> bool: + done_event = threading.Event() if wait else None + with self._lock: + self._spawn_xy = spawn_xy + if spawn_z is not None: + self._spawn_z = spawn_z + if spawn_yaw is not None: + self._spawn_yaw = spawn_yaw + self._reset_done_event = done_event + self._reset_requested.set() + if done_event is None: + return True + return done_event.wait(timeout) + + def ground_height_at( + self, + x: float, + y: float, + *, + ray_start_z: float = 10.0, + ) -> float | None: + """Raycast likely scene support geoms at ``x, y``. + + Scene collision geoms are emitted in group 3 and legacy floors are + commonly group 2. Restricting the ray to those groups avoids hitting + the robot itself when respawning near its current pose. + """ + with self._lock: + origin = np.array([float(x), float(y), float(ray_start_z)], dtype=np.float64) + direction = np.array([0.0, 0.0, -1.0], dtype=np.float64) + geom_id = np.array([-1], dtype=np.int32) + geomgroup = np.array([0, 0, 1, 1, 0, 0], dtype=np.uint8) + distance = mujoco.mj_ray( + self._model, + self._data, + origin, + direction, + geomgroup, + 1, + -1, + geom_id, + ) + if distance < 0.0: + return None + return float(ray_start_z - distance) + + def get_root_pose(self) -> tuple[NDArray[np.float64], NDArray[np.float64]] | None: + with self._lock: + return self.get_root_pose_unlocked() + + def get_root_pose_unlocked(self) -> tuple[NDArray[np.float64], NDArray[np.float64]] | None: + qpos_adr = self._root_qpos_adr + if qpos_adr is None: + return None + position = self._data.qpos[qpos_adr : qpos_adr + 3].copy() + qw, qx, qy, qz = self._data.qpos[qpos_adr + 3 : qpos_adr + 7].copy() + return position, np.array([qx, qy, qz, qw], dtype=np.float64) + def get_actuator_ctrl_range(self, joint_index: int) -> tuple[float, float] | None: mapping = self._joint_mappings[joint_index] if mapping.actuator_id is None: diff --git a/dimos/simulation/engines/mujoco_sim_module.py b/dimos/simulation/engines/mujoco_sim_module.py index ffc2effa39..82c4a60876 100644 --- a/dimos/simulation/engines/mujoco_sim_module.py +++ b/dimos/simulation/engines/mujoco_sim_module.py @@ -188,6 +188,12 @@ def post_step(self, engine: MujocoEngine) -> None: if self._gripper_idx < len(positions): shm.write_gripper_state(positions[self._gripper_idx]) + def clear_latched_commands(self) -> None: + self._latest_pd_pos_target = None + self._latest_pd_kp = None + self._latest_pd_kd = None + self._latest_pd_tau = None + def _gripper_joint_to_ctrl(self, joint_position: float) -> float: jlo, jhi = self._gripper_joint_range clo, chi = self._gripper_ctrl_range @@ -199,9 +205,24 @@ def _gripper_joint_to_ctrl(self, joint_position: float) -> float: class MujocoSimModuleConfig(ModuleConfig, DepthCameraConfig): - """Configuration for the unified MuJoCo simulation module.""" + """Configuration for the unified MuJoCo simulation module. + + Use ``address`` for the legacy path: one MJCF/MJB that already contains + scene and robot. Use ``robot_mjcf`` plus optional ``scene_xml`` for scene + packages: the module composes the scene wrapper, robot MJCF, and package + entities at startup. + """ address: str | Path = "" + scene_xml: str | Path | None = None + robot_mjcf: str | Path | None = None + robot_meshdir: str | Path | None = None + robot_id: str = "" + scene_entities: list[dict[str, Any]] = Field(default_factory=list) + spawn_xy: tuple[float, float] | None = None + spawn_z: float | None = None + spawn_yaw: float | None = None + reset_joint_positions: list[float] | None = None headless: bool = False dof: int = 7 @@ -294,6 +315,7 @@ def __init__(self, **kwargs: Any) -> None: # has a free joint at the robot root; None otherwise. self._imu_base_qpos_slice: slice | None = None self._root_base_qpos_adr: int | None = None + self._root_spawn_clearance_z: float | None = None @property def _camera_link(self) -> str: @@ -333,11 +355,15 @@ def get_depth_scale(self) -> float: @rpc def start(self) -> None: - if not self.config.address: - raise RuntimeError("MujocoSimModule: config.address (MJCF path) is required") + if not self.config.address and not self.config.robot_mjcf: + raise RuntimeError( + "MujocoSimModule: either config.address (legacy MJCF path) " + "or config.robot_mjcf (composed scene path) is required" + ) - # SHM key - adapter derives the same key from the same MJCF path. - shm_key = shm_key_from_path(self.config.address) + # SHM key - adapters derive the same key from the robot/control MJCF path. + shm_key_source = self.config.robot_mjcf or self.config.address + shm_key = shm_key_from_path(shm_key_source) self._shm = ManipShmWriter(shm_key) self._shm_ready_signaled = False @@ -372,13 +398,22 @@ def start(self) -> None: # Hooks are installed via set_step_hooks() after gripper detection # below, since they depend on the resolved gripper index. - self._engine = MujocoEngine( - config_path=Path(self.config.address), + engine_kwargs: dict[str, Any] = dict( headless=self.config.headless, cameras=cameras, - assets=engine_assets, robot_sim_spec=self.config.robot_sim_spec, + reset_joint_positions=self.config.reset_joint_positions, ) + if self.config.robot_mjcf is not None: + engine_kwargs["config_path"] = Path(self.config.robot_mjcf) + engine_kwargs["model"] = self._compose_model() + else: + engine_kwargs["config_path"] = Path(self.config.address) + engine_kwargs["assets"] = engine_assets + engine_kwargs["spawn_xy"] = self.config.spawn_xy + engine_kwargs["spawn_z"] = self.config.spawn_z + engine_kwargs["spawn_yaw"] = self.config.spawn_yaw + self._engine = MujocoEngine(**engine_kwargs) # Detect gripper (extra joint beyond dof). dof = self.config.dof @@ -423,6 +458,7 @@ def start(self) -> None: ) else: self._imu_base_qpos_slice = None + self._root_spawn_clearance_z = self._compute_root_spawn_clearance_z() # Wire SHM bridge hooks. self._sim_hooks = _WholeBodySimHooks( @@ -472,11 +508,46 @@ def start(self) -> None: logger.info( "MujocoSimModule started", address=self.config.address, + robot_mjcf=self.config.robot_mjcf, + scene_xml=self.config.scene_xml, dof=dof, camera=self.config.camera_name, shm_key=shm_key, ) + def _compose_model(self) -> mujoco.MjModel: + """Compose optional scene package MJCF + robot MJCF + entities.""" + from dimos.simulation.mujoco.entity_scene import add_entities_to_spec + + if self.config.robot_mjcf is None: + raise RuntimeError("MujocoSimModule: robot_mjcf is required for composition") + + if self.config.scene_xml is not None: + spec_scene = mujoco.MjSpec.from_file(str(self.config.scene_xml)) + else: + spec_scene = mujoco.MjSpec() + + spec_robot = mujoco.MjSpec.from_file(str(self.config.robot_mjcf)) + if self.config.robot_meshdir is not None: + spec_robot.meshdir = str(self.config.robot_meshdir) + + # Keep the robot controller timing stable when attached to a scene + # package whose wrapper may have different default options. + spec_scene.option.timestep = spec_robot.option.timestep + + spawn_xy = self.config.spawn_xy or (0.0, 0.0) + spawn_z = self.config.spawn_z if self.config.spawn_z is not None else 0.0 + frame = spec_scene.worldbody.add_frame( + pos=[float(spawn_xy[0]), float(spawn_xy[1]), float(spawn_z)], + ) + prefix = f"{self.config.robot_id}-" if self.config.robot_id else None + spec_scene.attach(spec_robot, prefix=prefix, frame=frame) + + if self.config.scene_entities: + add_entities_to_spec(spec_scene, self.config.scene_entities) + + return spec_scene.compile() + @rpc def stop(self) -> None: self._stop_event.set() @@ -509,6 +580,75 @@ def stop(self) -> None: op, err = errors[0] raise RuntimeError(f"MujocoSimModule.stop() failed during {op}: {err}") from err + @rpc + def reset(self) -> bool: + """Reset the running simulation to its current configured spawn pose.""" + engine = self._engine + if engine is None: + return False + if self._sim_hooks is not None: + self._sim_hooks.clear_latched_commands() + applied = engine.request_reset(wait=True) + logger.info("MujocoSimModule: reset requested", applied=applied) + return applied + + @rpc + def respawn_at( + self, + x: float, + y: float, + z: float | None = None, + yaw: float | None = None, + ) -> bool: + engine = self._engine + if engine is None: + return False + if self._sim_hooks is not None: + self._sim_hooks.clear_latched_commands() + + ground_z = None + spawn_z = None if z is None else float(z) + if spawn_z is None: + ground_z = engine.ground_height_at(float(x), float(y)) + if ground_z is not None and self._root_spawn_clearance_z is not None: + spawn_z = ground_z + self._root_spawn_clearance_z + + applied = engine.request_reset_to( + spawn_xy=(float(x), float(y)), + spawn_z=spawn_z, + spawn_yaw=None if yaw is None else float(yaw), + wait=True, + ) + logger.info( + "MujocoSimModule: respawn_at requested", + x=x, + y=y, + z=z, + ground_z=ground_z, + root_clearance_z=self._root_spawn_clearance_z, + spawn_z=spawn_z, + yaw=yaw, + applied=applied, + ) + return applied + + def _compute_root_spawn_clearance_z(self) -> float | None: + engine = self._engine + qpos_adr = self._root_base_qpos_adr + if engine is None or qpos_adr is None: + return None + + qpos = engine.data.qpos + root_x = float(qpos[qpos_adr]) + root_y = float(qpos[qpos_adr + 1]) + root_z = float(qpos[qpos_adr + 2]) + ground_z = engine.ground_height_at(root_x, root_y) + if ground_z is not None: + return root_z - ground_z + if self.config.spawn_z is not None: + return root_z - float(self.config.spawn_z) + return root_z + def _publish_shm_and_lcm(self, engine: MujocoEngine) -> None: """Post-step hook: SHM writes + LCM publishes. diff --git a/dimos/simulation/engines/test_mujoco_sim_module.py b/dimos/simulation/engines/test_mujoco_sim_module.py index 1a37941ef0..0074493858 100644 --- a/dimos/simulation/engines/test_mujoco_sim_module.py +++ b/dimos/simulation/engines/test_mujoco_sim_module.py @@ -14,12 +14,15 @@ from __future__ import annotations +from pathlib import Path from typing import Any from unittest.mock import MagicMock import numpy as np +import pytest -from dimos.simulation.engines.mujoco_sim_module import MujocoSimModule +from dimos.simulation.engines.mujoco_engine import MujocoEngine +from dimos.simulation.engines.mujoco_sim_module import MujocoSimModule, MujocoSimModuleConfig class _FakeData: @@ -32,6 +35,46 @@ class _FakeEngine: joint_names = ["joint_a", "joint_b"] +class _FakeRespawnEngine: + def __init__(self, *, ground_z: float = 0.08) -> None: + self.ground_z = ground_z + self.reset_requested = False + self.reset_to_kwargs: dict[str, Any] | None = None + + def request_reset(self, *, wait: bool) -> bool: + self.reset_requested = wait + return True + + def ground_height_at(self, x: float, y: float) -> float: + assert x == pytest.approx(2.6) + assert y == pytest.approx(0.0) + return self.ground_z + + def request_reset_to( + self, + *, + spawn_xy: tuple[float, float], + spawn_z: float | None, + spawn_yaw: float | None, + wait: bool, + ) -> bool: + self.reset_to_kwargs = { + "spawn_xy": spawn_xy, + "spawn_z": spawn_z, + "spawn_yaw": spawn_yaw, + "wait": wait, + } + return True + + +class _FakeSimHooks: + def __init__(self) -> None: + self.cleared = False + + def clear_latched_commands(self) -> None: + self.cleared = True + + def test_ready_signal_happens_after_joint_state_and_imu_write() -> None: events: list[str] = [] module = object.__new__(MujocoSimModule) @@ -63,3 +106,179 @@ def signal_ready(self, *, num_joints: int) -> None: module._publish_shm_and_lcm(_FakeEngine) assert events == ["joint_state", "imu", "ready"] + + +def test_reset_requests_engine_reset_and_clears_latched_commands() -> None: + module = object.__new__(MujocoSimModule) + engine = _FakeRespawnEngine() + hooks = _FakeSimHooks() + module._engine = engine + module._sim_hooks = hooks + + assert module.reset() is True + + assert engine.reset_requested is True + assert hooks.cleared is True + + +def test_respawn_at_uses_ground_height_plus_initial_root_clearance() -> None: + module = object.__new__(MujocoSimModule) + engine = _FakeRespawnEngine(ground_z=0.08) + hooks = _FakeSimHooks() + module._engine = engine + module._sim_hooks = hooks + module._root_spawn_clearance_z = 0.793 + + assert module.respawn_at(2.6, 0.0, yaw=0.25) is True + + assert engine.reset_to_kwargs == { + "spawn_xy": (2.6, 0.0), + "spawn_z": pytest.approx(0.873), + "spawn_yaw": 0.25, + "wait": True, + } + assert hooks.cleared is True + + +def _write_scene_xml(path: Path) -> None: + path.write_text( + """ + + +""".strip() + ) + + +def _write_robot_xml(path: Path) -> None: + path.write_text( + """ + + +""".strip() + ) + + +def _write_freejoint_xml(path: Path) -> None: + path.write_text( + """ + + +""".strip() + ) + + +def _scene_entity(entity_id: str) -> dict[str, object]: + return { + "id": entity_id, + "spawn": "initial", + "descriptor": { + "entity_id": entity_id, + "kind": "dynamic", + "mass": 1.0, + "shape_hint": "box", + "extents": [0.2, 0.2, 0.2], + }, + "initial_pose": { + "x": 1.0, + "y": 0.0, + "z": 0.1, + "qw": 1.0, + "qx": 0.0, + "qy": 0.0, + "qz": 0.0, + }, + } + + +@pytest.mark.mujoco +def test_compose_model_attaches_robot_before_scene_entities(tmp_path: Path) -> None: + mujoco = pytest.importorskip("mujoco") + scene_xml = tmp_path / "scene.xml" + robot_xml = tmp_path / "robot.xml" + _write_scene_xml(scene_xml) + _write_robot_xml(robot_xml) + + module = object.__new__(MujocoSimModule) + module.config = MujocoSimModuleConfig( + scene_xml=scene_xml, + robot_mjcf=robot_xml, + scene_entities=[_scene_entity("chair_000")], + spawn_xy=(0.25, -0.5), + spawn_z=0.8, + ) + + model = MujocoSimModule._compose_model(module) + + assert model.opt.timestep == pytest.approx(0.005) + assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "static_scene_box") >= 0 + assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "entity:chair_000") >= 0 + + free_joints: list[tuple[str, int]] = [] + for joint_id in range(model.njnt): + if int(model.jnt_type[joint_id]) != int(mujoco.mjtJoint.mjJNT_FREE): + continue + name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, joint_id) or "" + free_joints.append((name, int(model.jnt_qposadr[joint_id]))) + + robot_free = next(item for item in free_joints if item[0].endswith("floating_base_joint")) + entity_free = next(item for item in free_joints if item[0] == "entity:chair_000:free") + assert robot_free[1] == 0 + assert entity_free[1] > robot_free[1] + + engine = MujocoEngine(config_path=robot_xml, headless=True, model=model) + assert engine.model is model + assert engine.root_qpos_adr == 0 + assert any(name.endswith("hinge") for name in engine.joint_names) + + +@pytest.mark.mujoco +def test_engine_request_reset_to_applies_pose_in_sim_loop(tmp_path: Path) -> None: + robot_xml = tmp_path / "freejoint.xml" + _write_freejoint_xml(robot_xml) + + engine = MujocoEngine(config_path=robot_xml, headless=True) + assert engine.connect() is True + try: + assert engine.request_reset_to( + spawn_xy=(1.25, -0.5), + spawn_z=0.9, + spawn_yaw=0.3, + wait=True, + ) + pose = engine.get_root_pose() + assert pose is not None + position, quat_xyzw = pose + np.testing.assert_allclose(position, [1.25, -0.5, 0.9], atol=1e-8) + np.testing.assert_allclose( + quat_xyzw, + [0.0, 0.0, np.sin(0.15), np.cos(0.15)], + atol=1e-8, + ) + finally: + engine.disconnect() diff --git a/dimos/simulation/mujoco/entity_scene.py b/dimos/simulation/mujoco/entity_scene.py new file mode 100644 index 0000000000..c28c82c053 --- /dev/null +++ b/dimos/simulation/mujoco/entity_scene.py @@ -0,0 +1,347 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compose scene-package entities into a MuJoCo model via ``MjSpec``. + +The cook step removes entity prims (chairs, props) from the static +scene bake and writes their per-entity GLBs and metadata to the +package's ``entities/`` directory. At runtime, ``MujocoSimModule`` +attaches the robot first, then calls :func:`add_entities_to_spec` so +the robot keeps the first freejoint/qpos block and cooked entities +become first-class bodies after it in the composed model. + +Entities with ``kind == "dynamic"`` and positive mass receive a +freejoint (robot can push/grasp them); anything else is welded static. +Collision: primitive shapes (box/sphere/cylinder) use the descriptor +extents; mesh entities load the CoACD hulls cooked into the package +(``collision_paths`` in ``scene.meta.json``, written by +``dimos.experimental.pimsim.scene.entity_collision``). There is no +runtime decomposition — a mesh entity without cooked hulls falls back +to its AABB box with a warning to re-cook the package. + +A spawn-contact audit can be run on the compiled model to weld +entities that start in deep penetration with the static scene; see +:func:`spawn_penetrators`. The runtime caller (``MujocoSimModule``) +chooses when to invoke it. + +Body naming: ``entity:`` — consumers map MuJoCo bodies back +to entity ids through this prefix. +""" + +from __future__ import annotations + +import hashlib +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + import mujoco + + from dimos.simulation.scene_assets.spec import ScenePackage + +logger = setup_logger() + +ENTITY_BODY_PREFIX = "entity:" + +_MIN_HALF_EXTENT = 0.01 +# Sliding friction 0.3 ≈ furniture that scoots when bumped. Entity geoms +# carry priority=1 so this wins the contact pair outright (MuJoCo's +# default combine rule is element-wise max, which would let the μ=1.0 +# floor override it). Graspable props override via ``physics.friction``. +_DEFAULT_FRICTION = (0.3, 0.05, 0.001) +_DEFAULT_RGBA = (0.62, 0.62, 0.68, 1.0) +# Same geom group as the baked static scene so depth-based lidar renders +# (which hide robot groups 0/1) still see entities. +_ENTITY_GEOM_GROUP = 3 +# Spawn-contact audit: deeper penetration than this at t=0 demotes the +# entity to static instead of letting MuJoCo eject it. +_SPAWN_PENETRATION_LIMIT_M = 0.02 + + +def entity_body_name(entity_id: str) -> str: + return f"{ENTITY_BODY_PREFIX}{entity_id}" + + +def _initial_entities(entities: list[dict[str, Any]]) -> list[dict[str, Any]]: + return [e for e in entities if e.get("spawn", "initial") == "initial"] + + +def _box_size_and_offset(entity: dict[str, Any]) -> tuple[list[float], list[float]] | None: + """Half-extents + geom offset (body frame) for an entity's collision box.""" + descriptor = entity.get("descriptor", {}) + shape = descriptor.get("shape_hint", "mesh") + extents = [float(x) for x in descriptor.get("extents", [])] + + if shape == "box" and len(extents) == 3: + half = [max(x / 2.0, _MIN_HALF_EXTENT) for x in extents] + return half, [0.0, 0.0, 0.0] + if shape == "sphere" and len(extents) == 1: + r = max(extents[0], _MIN_HALF_EXTENT) + return [r, r, r], [0.0, 0.0, 0.0] + if shape == "cylinder" and len(extents) == 2: + r = max(extents[0], _MIN_HALF_EXTENT) + h = max(extents[1] / 2.0, _MIN_HALF_EXTENT) + return [r, r, h], [0.0, 0.0, 0.0] + + # Mesh entities: box from the cooked world-frame AABB. Cook poses are + # identity-rotation, so the AABB is axis-aligned in the body frame too. + aabb = entity.get("aabb") + pose = entity.get("initial_pose") + if not aabb or not pose: + return None + lo = [float(x) for x in aabb["min"]] + hi = [float(x) for x in aabb["max"]] + half = [max((h - low) / 2.0, _MIN_HALF_EXTENT) for low, h in zip(lo, hi, strict=True)] + center = [(h + low) / 2.0 for low, h in zip(lo, hi, strict=True)] + origin = [float(pose.get(k, 0.0)) for k in ("x", "y", "z")] + offset = [c - o for c, o in zip(center, origin, strict=True)] + return half, offset + + +def _entity_collision_hulls(entity: dict[str, Any]) -> list[Path]: + """Cooked collision hulls for a mesh entity, from package metadata. + + ``collision_paths`` is written by the cooker (CoACD decomposition of + the entity's visual GLB) and resolved to absolute paths by + ``load_scene_package``. The cooked per-entity GLBs are entity-local + (origin = initial_pose, world axes), so hulls drop in with zero geom + offset. All-or-nothing: a partially missing hull set falls back to + the AABB box rather than colliding with half a chair. + """ + raw = entity.get("collision_paths") + if not isinstance(raw, list) or not raw: + return [] + paths = [Path(item) for item in raw if isinstance(item, str) and item] + missing = [path for path in paths if not path.exists()] + if missing or not paths: + logger.warning( + "entity %s: %d/%d cooked collision hulls missing (e.g. %s); " + "falling back to AABB box — re-cook the scene package", + entity.get("id"), + len(missing), + len(paths), + missing[0] if missing else "", + ) + return [] + return paths + + +def _entity_friction(entity: dict[str, Any]) -> tuple[float, float, float]: + """``physics.friction`` from entity metadata (scalar sliding or full + [sliding, torsional, rolling] triple), else the scoot-able default.""" + raw = entity.get("physics", {}).get("friction") + sliding, torsional, rolling = _DEFAULT_FRICTION + if isinstance(raw, int | float): + sliding = float(raw) + elif isinstance(raw, list | tuple) and len(raw) == 3: + sliding, torsional, rolling = (float(v) for v in raw) + return sliding, torsional, rolling + + +def _entity_rgba(descriptor: dict[str, Any]) -> tuple[float, float, float, float]: + raw = descriptor.get("rgba") + if isinstance(raw, list | tuple) and len(raw) == 4: + return tuple(float(v) for v in raw) # type: ignore[return-value] + return _DEFAULT_RGBA + + +def add_entities_to_spec( + spec: mujoco.MjSpec, + entities: list[dict[str, Any]], + *, + force_static: frozenset[str] = frozenset(), +) -> None: + """Append scene-package entities as bodies on ``spec.worldbody``. + + Call after attaching the robot. Each ``spawn=="initial"`` entity + becomes one body named ``entity:`` with the descriptor's geom + shape and friction; dynamic entities also receive a freejoint named + ``entity::free``. + + ``force_static`` is a set of entity ids to demote to welded static + regardless of their ``kind`` — used by the spawn-penetration audit + when an entity boots in deep contact with the static scene. + """ + import mujoco + + for entity in _initial_entities(entities): + descriptor = entity.get("descriptor", {}) + entity_id = descriptor.get("entity_id") or entity.get("id") + pose = entity.get("initial_pose") + if not entity_id or not pose: + continue + entity_id = str(entity_id) + + kind = descriptor.get("kind", "kinematic") + mass = float(descriptor.get("mass", 0.0)) + dynamic = kind == "dynamic" and mass > 0.0 and entity_id not in force_static + + body = spec.worldbody.add_body( + name=entity_body_name(entity_id), + pos=[ + float(pose.get("x", 0.0)), + float(pose.get("y", 0.0)), + float(pose.get("z", 0.0)), + ], + quat=[ + float(pose.get("qw", 1.0)), + float(pose.get("qx", 0.0)), + float(pose.get("qy", 0.0)), + float(pose.get("qz", 0.0)), + ], + ) + if dynamic: + body.add_freejoint(name=f"{entity_body_name(entity_id)}:free") + + rgba = _entity_rgba(descriptor) + friction = _entity_friction(entity) + geom_kwargs: dict[str, Any] = dict( + name=f"{entity_body_name(entity_id)}:geom", + rgba=list(rgba), + friction=list(friction), + group=_ENTITY_GEOM_GROUP, + # priority=1: contact friction comes from the entity geom alone. + # MuJoCo's default combine rule (element-wise max across the + # pair) would otherwise let the μ=1.0 floor override every + # entity's friction. + priority=1, + ) + if dynamic: + geom_kwargs["mass"] = mass + + shape = descriptor.get("shape_hint", "mesh") + hull_paths = _entity_collision_hulls(entity) if shape == "mesh" else [] + if shape == "mesh" and not hull_paths and "collision_paths" not in entity: + logger.warning( + "entity %s: mesh entity has no cooked collision hulls; using AABB box " + "(re-cook the scene package with dimos.experimental.pimsim.scene.cook)", + entity_id, + ) + if hull_paths: + base_name = entity_body_name(entity_id) + if dynamic: + # MuJoCo derives per-geom mass from a body-level mass split + # across geoms by volume. Setting mass on each geom would + # double-count. Move mass to the body and drop it from the + # per-geom kwargs. + body.mass = mass + geom_kwargs.pop("mass", None) + for i, hull_obj in enumerate(hull_paths): + mesh_name = f"{base_name}:hull{i:03d}" + spec.add_mesh(name=mesh_name, file=str(hull_obj)) + # Name geoms uniquely so MuJoCo's compile doesn't reject + # duplicate-name collisions across multi-hull entities. + gk = dict(geom_kwargs) + gk["name"] = f"{base_name}:geom{i:03d}" + body.add_geom( + type=mujoco.mjtGeom.mjGEOM_MESH, + meshname=mesh_name, + **gk, + ) + continue + + box = _box_size_and_offset(entity) + if box is None: + logger.warning("entity %s has no usable collision shape; skipping", entity_id) + continue + half, offset = box + geom_kwargs["pos"] = offset + if shape == "sphere": + geom_type = mujoco.mjtGeom.mjGEOM_SPHERE + size = [half[0], 0.0, 0.0] + elif shape == "cylinder": + geom_type = mujoco.mjtGeom.mjGEOM_CYLINDER + size = [half[0], half[2], 0.0] + else: + geom_type = mujoco.mjtGeom.mjGEOM_BOX + size = half + body.add_geom(type=geom_type, size=size, **geom_kwargs) + + +def spawn_penetrators(model: mujoco.MjModel) -> frozenset[str]: + """Entity ids whose geoms start in deep contact at the spawn pose. + + Run after ``spec.compile()`` and before stepping; pass the returned + set as ``force_static`` to a second ``add_entities_to_spec`` call on + a fresh spec if you want to weld penetrating entities and recompile. + """ + import mujoco + + data = mujoco.MjData(model) + mujoco.mj_forward(model, data) + bad: set[str] = set() + for c in range(data.ncon): + contact = data.contact[c] + if contact.dist >= -_SPAWN_PENETRATION_LIMIT_M: + continue + for geom_id in (contact.geom1, contact.geom2): + name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, int(geom_id)) or "" + if name.startswith(ENTITY_BODY_PREFIX) and name.endswith(":geom"): + bad.add(name[len(ENTITY_BODY_PREFIX) : -len(":geom")]) + return frozenset(bad) + + +def compose_entity_model(scene_package: ScenePackage) -> Path | None: + """Legacy entry point — pre-compose a scene+entities ``.mjb`` from a + package wrapper. + + The new runtime path (``MujocoSimModule._compose_model``) calls + :func:`add_entities_to_spec` directly on the ``MjSpec`` and never + materialises a separate ``.mjb``. This function is retained for the + few callers that still want a precompiled binary; it returns ``None`` + when the package has no MuJoCo scene artifact. + """ + if scene_package.mujoco_scene_path is None: + return None + wrapper = Path(scene_package.mujoco_scene_path) + if not wrapper.exists(): + return None + + import mujoco + + entities = _initial_entities(scene_package.entities) + if not entities: + return wrapper + + spec = mujoco.MjSpec.from_file(str(wrapper)) + add_entities_to_spec(spec, entities) + model = spec.compile() + + penetrators = spawn_penetrators(model) + if penetrators: + logger.warning( + "%d entities spawn in deep contact and are welded static: %s", + len(penetrators), + ", ".join(sorted(penetrators)), + ) + spec = mujoco.MjSpec.from_file(str(wrapper)) + add_entities_to_spec(spec, entities, force_static=penetrators) + model = spec.compile() + + out_dir = wrapper.parent + key = hashlib.sha256(repr(entities).encode()).hexdigest()[:12] + mjb_path = out_dir / f"entities_{key}.mjb" + mujoco.mj_saveModel(model, str(mjb_path)) + return mjb_path + + +__all__ = [ + "ENTITY_BODY_PREFIX", + "add_entities_to_spec", + "compose_entity_model", + "entity_body_name", + "spawn_penetrators", +] diff --git a/dimos/simulation/scene_assets/spec.py b/dimos/simulation/scene_assets/spec.py new file mode 100644 index 0000000000..e24882ee48 --- /dev/null +++ b/dimos/simulation/scene_assets/spec.py @@ -0,0 +1,276 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Scene package contracts for offline asset cooking. + +Runtime modules consume the artifacts described here; they do not perform +the heavy bake themselves. +""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import asdict, dataclass, field +import json +from pathlib import Path +from typing import Any + +ARTIFACT_FRAMES = { + "browser_visual": "source", + "browser_collision": "source", + "mujoco": "dimos_world", +} + + +@dataclass(frozen=True) +class SceneMeshAlignment: + """Transform from a raw scene asset frame into DimOS world frame. + + Cookers use this to produce artifacts; runtime consumers keep it in + package metadata so viewers and simulators know each artifact's frame. + """ + + scale: float = 1.0 + translation: tuple[float, float, float] = (0.0, 0.0, 0.0) + rotation_zyx_deg: tuple[float, float, float] = (0.0, 0.0, 0.0) + y_up: bool = True + + +@dataclass(frozen=True) +class BrowserVisualSpec: + """Browser-rendered asset policy.""" + + enabled: bool = True + output_name: str = "visual.glb" + optimizer: str = "gltfpack" + simplify_ratio: float = 0.3 + simplify_error: float = 0.02 + texture_format: str | None = None + max_texture_size: int | None = None + max_meshes: int = 200 + max_materials: int = 50 + max_textures: int = 750 + max_vertices: int = 750_000 + max_vertex_growth_ratio: float = 1.25 + + +@dataclass(frozen=True) +class BrowserCollisionSpec: + """Browser raycast/physics collision asset policy.""" + + enabled: bool = True + output_name: str = "collision.glb" + target_faces: int = 100_000 + + +@dataclass(frozen=True) +class MujocoSceneSpec: + """MuJoCo collision asset policy.""" + + enabled: bool = True + include_visual_mesh: bool = False + + +@dataclass(frozen=True) +class SceneCookSpec: + """Complete cook input for one source scene.""" + + source_path: Path + alignment: SceneMeshAlignment = field(default_factory=SceneMeshAlignment) + browser_visual: BrowserVisualSpec = field(default_factory=BrowserVisualSpec) + browser_collision: BrowserCollisionSpec = field(default_factory=BrowserCollisionSpec) + mujoco: MujocoSceneSpec = field(default_factory=MujocoSceneSpec) + + +@dataclass(frozen=True) +class ScenePackage: + """Cooked scene outputs for runtime modules.""" + + package_dir: Path + source_path: Path + alignment: SceneMeshAlignment + visual_path: Path | None = None + browser_collision_path: Path | None = None + objects_path: Path | None = None + mujoco_scene_path: Path | None = None + metadata_path: Path | None = None + entities: list[dict[str, Any]] = field(default_factory=list) + stats: dict[str, Any] = field(default_factory=dict) + + def to_json_dict(self) -> dict[str, Any]: + package_dir = self.package_dir.expanduser().resolve() + + return { + "source_path": str(self.source_path), + "package_dir": ".", + "alignment": asdict(self.alignment), + "artifact_frames": ARTIFACT_FRAMES, + "artifacts": { + "browser_visual": _serialize_package_path(self.visual_path, package_dir), + "browser_collision": _serialize_package_path( + self.browser_collision_path, + package_dir, + ), + "objects": _serialize_package_path(self.objects_path, package_dir), + "mujoco_scene": _serialize_package_path(self.mujoco_scene_path, package_dir), + }, + "entities": _serialize_entity_paths(self.entities, package_dir), + "stats": self.stats, + } + + def write_metadata(self, path: Path | None = None) -> Path: + metadata_path = path or self.metadata_path or (self.package_dir / "scene.meta.json") + metadata_path.parent.mkdir(parents=True, exist_ok=True) + metadata_path.write_text(json.dumps(self.to_json_dict(), indent=2, sort_keys=True) + "\n") + return metadata_path + + +def load_scene_package(path: str | Path) -> ScenePackage: + """Load a previously written ``scene.meta.json``.""" + metadata_path = Path(path).expanduser().resolve() + raw = json.loads(metadata_path.read_text()) + _validate_artifact_frames(raw, metadata_path) + artifacts = raw.get("artifacts", {}) + align = SceneMeshAlignment(**raw["alignment"]) + package_dir = _resolve_package_dir(raw.get("package_dir"), metadata_path) + return ScenePackage( + package_dir=package_dir, + source_path=Path(raw["source_path"]), + alignment=align, + visual_path=_resolve_package_path(artifacts.get("browser_visual"), package_dir), + browser_collision_path=( + _resolve_package_path(artifacts.get("browser_collision"), package_dir) + ), + objects_path=_resolve_package_path(artifacts.get("objects"), package_dir), + mujoco_scene_path=_resolve_package_path(artifacts.get("mujoco_scene"), package_dir), + metadata_path=metadata_path, + entities=_resolve_entity_paths(raw.get("entities", []), package_dir), + stats=raw.get("stats", {}), + ) + + +def _validate_artifact_frames(raw: dict[str, Any], metadata_path: Path) -> None: + frames = raw.get("artifact_frames") + if frames is None: + raise ValueError( + f"scene package is missing artifact frame metadata: {metadata_path}. " + "Recook it with dimos.experimental.pimsim.scene.cook." + ) + + artifacts = raw.get("artifacts", {}) + required = { + "browser_visual": "browser_visual", + "browser_collision": "browser_collision", + "mujoco_scene": "mujoco", + } + for artifact_name, frame_name in required.items(): + if artifacts.get(artifact_name) and frames.get(frame_name) != ARTIFACT_FRAMES[frame_name]: + raise ValueError( + f"scene package artifact frame mismatch in {metadata_path}: " + f"{frame_name}={frames.get(frame_name)!r}, " + f"expected {ARTIFACT_FRAMES[frame_name]!r}. Recook the scene package." + ) + + +def _serialize_package_path(path: Path | None, package_dir: Path) -> str | None: + if path is None: + return None + resolved = path.expanduser().resolve() + try: + return resolved.relative_to(package_dir).as_posix() + except ValueError: + return str(resolved) + + +def _resolve_package_dir(raw: str | None, metadata_path: Path) -> Path: + if raw is None: + return metadata_path.parent + path = Path(raw).expanduser() + if path.is_absolute(): + return path + return (metadata_path.parent / path).resolve() + + +def _resolve_package_path(raw: str | None, package_dir: Path) -> Path | None: + if not raw: + return None + path = Path(raw).expanduser() + if path.is_absolute(): + return path + return (package_dir / path).resolve() + + +_ENTITY_PATH_KEYS = ("visual_path", "collision_path", "mesh_path") +_ENTITY_PATH_LIST_KEYS = ("collision_paths",) + + +def _serialize_entity_paths( + entities: list[dict[str, Any]], package_dir: Path +) -> list[dict[str, Any]]: + out = deepcopy(entities) + for entity in out: + _rewrite_entity_paths(entity, package_dir, serialize=True) + return out + + +def _resolve_entity_paths( + entities: list[dict[str, Any]], package_dir: Path +) -> list[dict[str, Any]]: + out = deepcopy(entities) + for entity in out: + _rewrite_entity_paths(entity, package_dir, serialize=False) + return out + + +def _rewrite_entity_paths( + entity: dict[str, Any], + package_dir: Path, + *, + serialize: bool, +) -> None: + def rewrite(value: Any) -> Any: + if not isinstance(value, str): + return value + path = Path(value).expanduser() + if serialize: + return _serialize_package_path(path, package_dir) + if path.is_absolute(): + return str(path) + return str((package_dir / path).resolve()) + + for key in _ENTITY_PATH_KEYS: + if key in entity: + entity[key] = rewrite(entity[key]) + + for key in _ENTITY_PATH_LIST_KEYS: + value = entity.get(key) + if isinstance(value, list): + entity[key] = [rewrite(item) for item in value] + + artifacts = entity.get("artifacts") + if isinstance(artifacts, dict): + for key, value in list(artifacts.items()): + artifacts[key] = rewrite(value) + + +__all__ = [ + "ARTIFACT_FRAMES", + "BrowserCollisionSpec", + "BrowserVisualSpec", + "MujocoSceneSpec", + "SceneCookSpec", + "SceneMeshAlignment", + "ScenePackage", + "load_scene_package", +] diff --git a/dimos/simulation/scenes/catalog.py b/dimos/simulation/scenes/catalog.py new file mode 100644 index 0000000000..13f56100e8 --- /dev/null +++ b/dimos/simulation/scenes/catalog.py @@ -0,0 +1,134 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Named scene packages for simulation/runtime viewers.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from dimos.simulation.scene_assets.spec import ScenePackage, load_scene_package +from dimos.utils.data import get_data + +DEFAULT_SCENE = "dimos-office" +_DISABLED_SCENE_NAMES = {"", "none", "off", "disabled", "false", "0"} +_ALIASES = { + "default": DEFAULT_SCENE, + "office": DEFAULT_SCENE, + "dimos-office": DEFAULT_SCENE, + "dimos_office": DEFAULT_SCENE, + "office-splat": "dimos-office-splat", + "dimos-office-splat": "dimos-office-splat", + "dimos_office_splat": "dimos-office-splat", + "street": "street-lite", + "street-lite": "street-lite", + "street_lite": "street-lite", + "mall": "mall-babylon-nolights", + "mall_babylon_nolights": "mall-babylon-nolights", + "mall-babylon-nolights": "mall-babylon-nolights", + "lowpoly": "lowpoly-tdm", + "lowpoly-tdm": "lowpoly-tdm", + "lowpoly_tdm": "lowpoly-tdm", + "tdm": "lowpoly-tdm", + "apartment": "dimos-apartment", + "dimos-apartment": "dimos-apartment", + "dimos_apartment": "dimos-apartment", +} +_PACKAGE_DIRS = { + DEFAULT_SCENE: "dimos_office", + "dimos-office-splat": "dimos_office_splat", + "street-lite": "street_lite", + "mall-babylon-nolights": "mall_babylon_nolights", + "lowpoly-tdm": "lowpoly_tdm", + "dimos-apartment": "dimos_apartment", +} + + +def resolve_scene_package( + scene: str | Path | None = None, + **_legacy: Any, +) -> ScenePackage | None: + """Resolve a scene name, metadata path, or package directory. + + ``robot_mjcf_path`` and ``meshdir`` are accepted as legacy keyword + arguments and ignored; scene packages are now robot-agnostic and the + runtime composer attaches the robot via ``MjSpec.attach()``. + """ + if scene is None: + scene = DEFAULT_SCENE + + scene_text = str(scene).strip() + if scene_text.lower() in _DISABLED_SCENE_NAMES: + return None + + candidate = Path(scene_text).expanduser() + if candidate.exists(): + if candidate.is_dir(): + metadata_path = candidate / "scene.meta.json" + if not metadata_path.exists(): + raise FileNotFoundError(f"scene package directory has no {metadata_path.name}") + elif candidate.name == "scene.meta.json" or candidate.suffix.lower() == ".json": + metadata_path = candidate + else: + raise ValueError( + "scene paths must point to a cooked scene.meta.json or package directory; " + f"got raw asset path: {candidate}" + ) + return load_scene_package(metadata_path) + + name = _ALIASES.get(scene_text.lower()) + if name is None: + known = ", ".join(sorted(_PACKAGE_DIRS)) + raise ValueError(f"unknown scene '{scene_text}'. Known scenes: {known}") + + if name == DEFAULT_SCENE: + return _resolve_dimos_office() + + metadata_path = _scene_package_dir() / _PACKAGE_DIRS[name] / "scene.meta.json" + if not metadata_path.exists(): + raise FileNotFoundError(f"scene package '{name}' is not cooked yet: {metadata_path}") + return load_scene_package(metadata_path) + + +def _resolve_dimos_office() -> ScenePackage: + metadata_path = _scene_package_dir() / _PACKAGE_DIRS[DEFAULT_SCENE] / "scene.meta.json" + if not metadata_path.exists(): + raise FileNotFoundError( + "dimos-office scene package is not cooked yet: " + f"{metadata_path}. Run dimos.experimental.pimsim.scene.cook first." + ) + + package = load_scene_package(metadata_path) + if ( + package.visual_path is not None + and package.browser_collision_path is not None + and package.mujoco_scene_path is not None + and package.visual_path.exists() + and package.browser_collision_path.exists() + and package.mujoco_scene_path.exists() + ): + return package + + raise ValueError( + "dimos-office scene package is incomplete or has missing artifacts: " + f"{metadata_path}. Recook it from the bundled office scene." + ) + + +def _scene_package_dir() -> Path: + return get_data("scene_packages") + + +__all__ = ["DEFAULT_SCENE", "resolve_scene_package"] diff --git a/dimos/simulation/utils/xml_parser.py b/dimos/simulation/utils/xml_parser.py index 052657ea95..fb3858be9c 100644 --- a/dimos/simulation/utils/xml_parser.py +++ b/dimos/simulation/utils/xml_parser.py @@ -44,8 +44,10 @@ class _ActuatorSpec: tendon: str | None -def build_joint_mappings(xml_path: Path, model: mujoco.MjModel) -> list[JointMapping]: - specs = _parse_actuator_specs(xml_path) +def build_joint_mappings(xml_path: Path | None, model: mujoco.MjModel) -> list[JointMapping]: + specs: list[_ActuatorSpec] = [] + if xml_path is not None: + specs = _parse_actuator_specs(xml_path) if specs: return _build_joint_mappings_from_specs(specs, model) if int(model.nu) > 0: diff --git a/stubs/mujoco/__init__.pyi b/stubs/mujoco/__init__.pyi index ae4ddbf162..9785de546b 100644 --- a/stubs/mujoco/__init__.pyi +++ b/stubs/mujoco/__init__.pyi @@ -35,6 +35,19 @@ class MjData: def __getattr__(self, name: str) -> Any: ... def __init__(self, model: MjModel) -> None: ... +class MjSpec: + meshdir: str + worldbody: Any + # The MjSpec editing API is large (add_mesh, add_body, attach, …); rare + # accesses are Any, same approach as MjModel. + def __getattr__(self, name: str) -> Any: ... + def __init__(self) -> None: ... + @classmethod + def from_file(cls, path: str) -> MjSpec: ... + @classmethod + def from_string(cls, xml: str) -> MjSpec: ... + def compile(self) -> MjModel: ... + class MjvOption: def __init__(self) -> None: ... @@ -57,6 +70,7 @@ def mj_step(model: MjModel, data: MjData, nstep: int = ...) -> None: ... def mj_resetDataKeyframe(model: MjModel, data: MjData, key: int) -> None: ... def mj_name2id(model: MjModel, type: int, name: str | None) -> int: ... def mj_id2name(model: MjModel, type: int, id: int) -> str | None: ... +def mj_saveModel(model: MjModel, filename: str, buffer: Any = ...) -> None: ... def set_mjcb_control( cb: Callable[[MjModel, MjData], None] | None, ) -> None: ... @@ -67,9 +81,20 @@ class mjtObj: mjOBJ_ACTUATOR: int mjOBJ_BODY: int mjOBJ_CAMERA: int + mjOBJ_GEOM: int mjOBJ_JOINT: int + mjOBJ_MESH: int + mjOBJ_SITE: int mjOBJ_TENDON: int +class mjtGeom: + mjGEOM_PLANE: int + mjGEOM_SPHERE: int + mjGEOM_CAPSULE: int + mjGEOM_CYLINDER: int + mjGEOM_BOX: int + mjGEOM_MESH: int + class mjtJoint: mjJNT_HINGE: int mjJNT_SLIDE: int From fc150e30cf16aa1b528c52173e8d37e354b0d297 Mon Sep 17 00:00:00 2001 From: Nabla7 Date: Thu, 25 Jun 2026 02:17:42 +0200 Subject: [PATCH 02/73] Add scene package cooking pipeline --- data/.lfs/dimos_office_mesh.tar.gz | 3 + dimos/experimental/pimsim/scene/README.md | 312 +++++ .../pimsim/scene/browser_collision.py | 249 ++++ dimos/experimental/pimsim/scene/cook.py | 330 ++++++ .../pimsim/scene/entity_collision.py | 141 +++ dimos/experimental/pimsim/scene/inspect.py | 185 +++ dimos/experimental/pimsim/scene/plan.py | 341 ++++++ dimos/experimental/pimsim/scene/sidecar.py | 189 ++++ .../pimsim/scene/test_entity_collision.py | 67 ++ dimos/experimental/pimsim/scene/test_spec.py | 324 ++++++ .../pimsim/scene/visual_blender.py | 313 +++++ dimos/experimental/pimsim/scene/visual_glb.py | 363 ++++++ dimos/simulation/mujoco/collision_spec.py | 968 ++++++++++++++++ dimos/simulation/mujoco/scene_mesh_to_mjcf.py | 1003 +++++++++++++++++ .../simulation/mujoco/test_collision_spec.py | 65 ++ dimos/simulation/scene_assets/mesh_scene.py | 710 ++++++++++++ pyproject.toml | 19 +- uv.lock | 109 +- 18 files changed, 5663 insertions(+), 28 deletions(-) create mode 100644 data/.lfs/dimos_office_mesh.tar.gz create mode 100644 dimos/experimental/pimsim/scene/README.md create mode 100644 dimos/experimental/pimsim/scene/browser_collision.py create mode 100644 dimos/experimental/pimsim/scene/cook.py create mode 100644 dimos/experimental/pimsim/scene/entity_collision.py create mode 100644 dimos/experimental/pimsim/scene/inspect.py create mode 100644 dimos/experimental/pimsim/scene/plan.py create mode 100644 dimos/experimental/pimsim/scene/sidecar.py create mode 100644 dimos/experimental/pimsim/scene/test_entity_collision.py create mode 100644 dimos/experimental/pimsim/scene/test_spec.py create mode 100644 dimos/experimental/pimsim/scene/visual_blender.py create mode 100644 dimos/experimental/pimsim/scene/visual_glb.py create mode 100644 dimos/simulation/mujoco/collision_spec.py create mode 100644 dimos/simulation/mujoco/scene_mesh_to_mjcf.py create mode 100644 dimos/simulation/mujoco/test_collision_spec.py create mode 100644 dimos/simulation/scene_assets/mesh_scene.py diff --git a/data/.lfs/dimos_office_mesh.tar.gz b/data/.lfs/dimos_office_mesh.tar.gz new file mode 100644 index 0000000000..d9a02c20b2 --- /dev/null +++ b/data/.lfs/dimos_office_mesh.tar.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7725fd12a8f88e8f757478d3afe9a06938a21d718737c80b9655c9725813c1bf +size 318485597 diff --git a/dimos/experimental/pimsim/scene/README.md b/dimos/experimental/pimsim/scene/README.md new file mode 100644 index 0000000000..a2e52079cb --- /dev/null +++ b/dimos/experimental/pimsim/scene/README.md @@ -0,0 +1,312 @@ +# Scene Packages + +A scene package is the cooked, robot-agnostic form of a 3D environment. It +contains the visual mesh, collision geometry, per-object metadata, a scene-only +MuJoCo wrapper, and optional runtime entities. + +The runtime rule is simple: + +```text +source asset + sidecar -> cooked scene package -> simulator attaches any robot +``` + +The robot is not part of the package. The cooker prepares the world once, +offline. `MujocoSimModule` loads the cooked world at runtime and attaches the +robot MJCF into the same `MjSpec`. + +## Package Layout + +```text +data/scene_packages// +├── scene.meta.json manifest: alignment, artifact paths, entities, stats +├── mujoco// +│ ├── wrapper.xml scene-only MJCF, no robot +│ └── *.obj static collision assets +├── entities// +│ ├── visual.glb per-entity visual, in entity-local frame +│ └── mujoco_collision/ cook-time convex hulls +└── browser/ + ├── visual.glb + ├── collision.glb + └── objects.json semantic object table for browser/raycast users +``` + +Packages are content-hash keyed on the source mesh, alignment, sidecar, and cook +schema version. Changing one of those inputs creates a new cooked output. + +## Spec And Backends + +`dimos/simulation/scene_assets/spec.py` is the shared scene-package contract. It +defines: + +- cook input specs such as `SceneCookSpec`, `BrowserVisualSpec`, + `BrowserCollisionSpec`, and `MujocoSceneSpec`; +- the runtime `ScenePackage` object; +- the on-disk `scene.meta.json` shape; +- `load_scene_package()`, which resolves package-relative artifact paths and + validates artifact frame metadata. + +This spec is intentionally general. A module or simulator that wants to consume +a scene package should not parse `scene.meta.json` by hand. It should load the +package once: + +```python +from dimos.simulation.scene_assets.spec import load_scene_package + +package = load_scene_package("data/scene_packages/dimos_office/scene.meta.json") +``` + +and then consume the artifact it understands: + +- browser/viewer systems use `package.visual_path`, `package.browser_collision_path`, + and `package.objects_path`; +- MuJoCo uses `package.mujoco_scene_path` plus `package.entities`; +- future simulators can add their own artifact fields/specs without changing the + source mesh loader or the scene-name catalog. + +`dimos/simulation/scene_assets/mesh_scene.py` is a different layer. It loads raw +source geometry and applies `SceneMeshAlignment`; it is used by cook-time tools, +not by normal runtime modules. Runtime modules should prefer the cooked package +contract from `spec.py`. + +`dimos/simulation/scenes/catalog.py` is only name/path resolution. It maps user +inputs like `--scene office` to a loaded `ScenePackage`. It does not define what +a package is, and it should not know how MuJoCo, browser viewers, or other +systems load their artifacts. + +## Workflow + +1. Inspect the source asset and identify the prims that need authored collision. +2. Write `.cook.json` next to the source mesh. +3. Cook the package. +4. Verify the generated `mujoco//wrapper.xml`. +5. Load it through a blueprint with `--scene`. + +The DimOS office scene is the reference example below. + +## Office Example + +The office source mesh lives in data: + +```text +data/dimos_office_mesh/dimos_office_mesh.glb +``` + +That file is a visual asset, not a physics contract. Before writing the sidecar, +inspect the scene graph and bounds of the source prims after applying the same +alignment used by the cooker: + +```bash +python - <<'PY' +from pathlib import Path + +import numpy as np + +from dimos.simulation.scene_assets.mesh_scene import ( + SceneMeshAlignment, + load_scene_prims, +) + +source = Path("data/dimos_office_mesh/dimos_office_mesh.glb") +alignment = SceneMeshAlignment(scale=2.0, y_up=False) + +for prim in load_scene_prims(source, alignment=alignment): + name = prim.visual_node_name or prim.prim_path or prim.name + if "Floor" not in name: + continue + lo = np.min(prim.vertices, axis=0) + hi = np.max(prim.vertices, axis=0) + extent = hi - lo + print(f"{name}: min={lo.round(4).tolist()} max={hi.round(4).tolist()} extent={extent.round(4).tolist()}") +PY +``` + +For the current office GLB, the relevant support prim is `Floor_Plane.002`. Its +Z extent is effectively zero: it is a visual sheet. If that sheet is cooked into +a very thin MuJoCo box, humanoid foot contacts can clip through it during walking +even though the robot can stand still on it. + +Do not solve that by adding an infinite MuJoCo plane. That ignores the authored +scene geometry and breaks as soon as floors have different heights, holes, +stairs, ramps, platforms, or multiple stories. The sidecar should instead state +the authored physics intent for the named support surface. + +Create or edit: + +```text +data/dimos_office_mesh/dimos_office_mesh.cook.json +``` + +with the floor override: + +```json +{ + "collision": { + "default": "auto", + "prim_overrides": { + "Floor*": { + "type": "box", + "min_thickness": 0.04, + "preserve": "top" + } + } + } +} +``` + +Reasoning: + +- `Floor*` matches `Floor_Plane.002` without depending on the sanitized MJCF + name. +- `type: "box"` says the floor is a support slab, not a raw visual mesh. +- `min_thickness: 0.04` gives the support 4 cm of total thickness. +- `preserve: "top"` keeps the walkable surface at the authored visual height and + expands the slab downward. + +Cook the office package: + +```bash +python -m dimos.experimental.pimsim.scene.cook \ + data/dimos_office_mesh/dimos_office_mesh.glb \ + --cook-spec data/dimos_office_mesh/dimos_office_mesh.cook.json \ + --output-dir data/scene_packages/dimos_office \ + --scale 2.0 \ + --no-y-up \ + --rebake +``` + +Verify the generated MuJoCo wrapper: + +```bash +rg 'Floor.*_col' data/scene_packages/dimos_office/mujoco -g wrapper.xml +``` + +For a floor whose top is at `z=0`, the cooked geom should have half-thickness +`0.02` and center `z=-0.02`, for example: + +```xml + +``` + +The important invariant is top height unchanged, bottom expanded downward: + +```text +top_z = pos_z + size_z +bottom_z = pos_z - size_z +``` + +## Loading In MuJoCo + +The G1 GR00T WBC blueprint already uses the scene package path when `--scene` +is provided: + +```bash +python -m dimos.robot.cli.dimos \ + --simulation mujoco \ + --scene office \ + --viewer rerun \ + run unitree-g1-groot-wbc +``` + +`office` resolves through `dimos/simulation/scenes/catalog.py` to: + +```text +data/scene_packages/dimos_office/scene.meta.json +``` + +At runtime, `MujocoSimModule`: + +1. receives a `ScenePackage` resolved by `catalog.py`; +2. reads `package.mujoco_scene_path`, which points at + `mujoco//wrapper.xml`; +3. loads the robot-only G1 MJCF; +4. attaches the robot into the scene `MjSpec`; +5. adds `package.entities`; +6. compiles one MuJoCo model. + +The robot MJCF must stay robot-only: no office floor, no scene walls, no +furniture, no manipulation rig. Scene geometry belongs in the cooked package. + +## Sidecar Schema + +`.cook.json` can contain static-collision policy and interactables: + +```json +{ + "collision": { + "default": "auto", + "prim_overrides": { + "Floor*": {"type": "box", "min_thickness": 0.04, "preserve": "top"}, + "Wall_*": {"type": "box"}, + "Stairs_*": {"type": "decompose", "max_hulls": 16} + } + }, + "interactables": [ + { + "id": "chair_000", + "source_prim_paths": ["Chair_000_*"], + "mass": 8.0, + "physics": {"shape": "mesh"}, + "tags": ["chair"] + } + ] +} +``` + +Static collision types are: + +```text +auto | box | sphere | cylinder | capsule | plane | hull | mesh | decompose | skip +``` + +For box overrides, `min_thickness` is full world-Z thickness in meters. +`preserve` can be `top`, `bottom`, or `center`. + +Interactables become MuJoCo bodies named `entity:`. They can be: + +- extracted from the source mesh with `source_prim_paths`; +- synthetic primitives with an authored `pose` and `physics.extents`; +- static, kinematic, or dynamic depending on `kind` and `mass`. + +## Data Publishing + +Scene source meshes and cooked packages live under `data/`, which is ignored by +normal git. Do not add these artifacts with ordinary `git add`. Publish or update +them through the repository LFS bin workflow described in: + +```text +docs/development/large_file_management.md +``` + +Code and docs changes can go through normal git. Data changes such as +`data/dimos_office_mesh/dimos_office_mesh.cook.json` and +`data/scene_packages/dimos_office` should be handled through the LFS script when +we are ready to ship them. + +## Why Bake + +MuJoCo treats a mesh geom as convex for collision. A raw concave building mesh is +therefore the wrong collision representation: it becomes a coarse blob or a bad +support surface. The cooker turns source prims into MuJoCo-friendly collision: +primitives where the shape is obvious, hulls where they fit, and CoACD +decompositions for concave objects. + +The sidecar is the place for source-specific knowledge. It keeps the cook +deterministic and reviewable instead of relying on broad heuristics that may fail +on multi-story buildings or unusual floor geometry. + +## Reference + +| File | Role | +|---|---| +| `dimos/experimental/pimsim/scene/cook.py` | cook entry point and CLI | +| `dimos/experimental/pimsim/scene/sidecar.py` | `.cook.json` schema | +| `dimos/experimental/pimsim/scene/plan.py` | sidecar to cook plan | +| `dimos/experimental/pimsim/scene/inspect.py` | source asset statistics | +| `dimos/simulation/scene_assets/spec.py` | shared scene-package metadata contract | +| `dimos/simulation/scene_assets/mesh_scene.py` | source mesh loading and prim inspection | +| `dimos/simulation/mujoco/scene_mesh_to_mjcf.py` | MuJoCo scene bake | +| `dimos/simulation/mujoco/collision_spec.py` | static collision policy | +| `dimos/simulation/mujoco/entity_scene.py` | runtime entity composition | +| `dimos/simulation/scenes/catalog.py` | scene name/path resolution | +| `dimos/simulation/engines/mujoco_sim_module.py` | runtime scene + robot composition | diff --git a/dimos/experimental/pimsim/scene/browser_collision.py b/dimos/experimental/pimsim/scene/browser_collision.py new file mode 100644 index 0000000000..0b7e9f35dd --- /dev/null +++ b/dimos/experimental/pimsim/scene/browser_collision.py @@ -0,0 +1,249 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Bake browser-side collision geometry from a scene asset.""" + +from __future__ import annotations + +from dataclasses import dataclass +import json +from pathlib import Path +from typing import Any + +import numpy as np +import open3d as o3d # type: ignore[import-untyped] +import trimesh + +from dimos.experimental.pimsim.scene.inspect import inspect_scene_asset +from dimos.simulation.mujoco.collision_spec import CollisionSpec +from dimos.simulation.scene_assets.mesh_scene import ( + SceneMeshAlignment, + ScenePrimMesh, + load_scene_prims, + split_disconnected_scene_prims, +) +from dimos.simulation.scene_assets.spec import BrowserCollisionSpec +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +OBJECTS_SIDECAR_NAME = "objects.json" + + +@dataclass(frozen=True) +class BrowserCollisionCookResult: + path: Path + stats: dict[str, Any] + objects_path: Path | None = None + + +def cook_browser_collision( + source_path: str | Path, + output_dir: str | Path, + *, + alignment: SceneMeshAlignment | None = None, + spec: BrowserCollisionSpec | None = None, + collision_spec: CollisionSpec | None = None, + rebake: bool = False, +) -> BrowserCollisionCookResult | None: + """Write a simplified GLB used for browser picking/raycast/physics. + + For scene packages this should stay in source-asset coordinates; the + browser applies the package alignment to the visual and collision roots + together. + """ + browser_spec = spec or BrowserCollisionSpec() + if not browser_spec.enabled: + return None + + source = Path(source_path).expanduser().resolve() + out_dir = Path(output_dir).expanduser().resolve() + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / browser_spec.output_name + objects_path = out_dir / OBJECTS_SIDECAR_NAME + + mesh_cached = out_path.exists() and not rebake + objects_cached = objects_path.exists() and not rebake + if mesh_cached and objects_cached: + return BrowserCollisionCookResult( + path=out_path, + stats=inspect_scene_asset(out_path).to_json_dict(), + objects_path=objects_path, + ) + + prims = _load_collision_prims(source, alignment=alignment, collision_spec=collision_spec) + stats: dict[str, Any] + if mesh_cached: + stats = inspect_scene_asset(out_path).to_json_dict() + else: + mesh = _build_fused_collision_mesh( + prims, collision_spec or CollisionSpec.auto_discover(source) + ) + original_triangles = len(mesh.triangles) + target_faces = int(browser_spec.target_faces) + if target_faces > 0 and original_triangles > target_faces: + logger.info( + "browser collision: simplifying %s triangles -> %s", + original_triangles, + target_faces, + ) + mesh = mesh.simplify_quadric_decimation(target_number_of_triangles=target_faces) + mesh.remove_degenerate_triangles() + mesh.remove_duplicated_triangles() + mesh.remove_duplicated_vertices() + mesh.remove_non_manifold_edges() + _write_glb(mesh, out_path) + stats = inspect_scene_asset(out_path).to_json_dict() + stats["source_triangles"] = original_triangles + stats["target_faces"] = target_faces + + objects = extract_scene_objects(prims) + if not objects_cached: + _write_objects_json(objects_path, objects) + stats["objects"] = len(objects) + return BrowserCollisionCookResult(path=out_path, stats=stats, objects_path=objects_path) + + +def _write_glb(mesh: o3d.geometry.TriangleMesh, path: Path) -> None: + vertices = np.asarray(mesh.vertices, dtype=np.float64) + faces = np.asarray(mesh.triangles, dtype=np.int64) + if len(vertices) == 0 or len(faces) == 0: + raise RuntimeError("browser collision bake produced an empty mesh") + trimesh.Trimesh(vertices=vertices, faces=faces, process=False).export(str(path)) + + +def _load_collision_prims( + source: Path, + *, + alignment: SceneMeshAlignment | None, + collision_spec: CollisionSpec | None, +) -> list[ScenePrimMesh]: + spec = collision_spec or CollisionSpec.auto_discover(source) + source_alignment = alignment or SceneMeshAlignment(y_up=False) + + prims = load_scene_prims(source, alignment=source_alignment) + if spec.split_disconnected_components: + prims, split_stats = split_disconnected_scene_prims( + prims, + min_components=spec.split_min_components, + extent_ratio=spec.split_extent_ratio, + prim_min_extent=spec.split_prim_min_extent_m, + axis_ratio=spec.split_axis_ratio, + min_component_extent=spec.split_component_min_extent_m, + min_component_faces=spec.split_component_min_faces, + can_split=lambda prim: ( + spec.resolve(prim.prim_path or prim.name).get("type", spec.default) == "auto" + ), + ) + if split_stats["split_prims"]: + logger.info( + "browser collision: split %s disconnected prims into %s kept " + "components; dropped %s tiny components", + split_stats["split_prims"], + split_stats["emitted_components"], + split_stats["dropped_components"], + ) + return prims + + +def _build_fused_collision_mesh( + prims: list[ScenePrimMesh], + spec: CollisionSpec, +) -> o3d.geometry.TriangleMesh: + vertices: list[np.ndarray] = [] + faces: list[np.ndarray] = [] + vertex_offset = 0 + for prim in prims: + mesh = _mesh_for_prim(prim, spec) + if mesh is None: + continue + prim_vertices = np.asarray(mesh.vertices, dtype=np.float64) + prim_faces = np.asarray(mesh.triangles, dtype=np.int64) + if len(prim_vertices) == 0 or len(prim_faces) == 0: + continue + vertices.append(prim_vertices) + faces.append(prim_faces + vertex_offset) + vertex_offset += len(prim_vertices) + if not vertices: + raise RuntimeError("browser collision sidecar skipped every prim") + return _mesh_from_arrays(np.concatenate(vertices, axis=0), np.concatenate(faces, axis=0)) + + +def extract_scene_objects(prims: list[ScenePrimMesh]) -> list[dict[str, Any]]: + """Per-prim semantic metadata (id, prim_path, AABB in source frame). + + Emitted independently of the fused collision GLB so the runtime can + answer ``findAsset("sectional")``-style queries without paying a + per-object PhysicsAggregate cost. AABB shares the collision GLB + frame (source / z-up after alignment). + """ + objects: list[dict[str, Any]] = [] + for prim in prims: + v = np.asarray(prim.vertices, dtype=np.float64) + if v.size == 0: + continue + objects.append( + { + "id": prim.name, + "prim_path": prim.prim_path, + "aabb": { + "min": v.min(axis=0).tolist(), + "max": v.max(axis=0).tolist(), + }, + } + ) + return objects + + +def _write_objects_json(path: Path, objects: list[dict[str, Any]]) -> None: + payload = {"frame": "source", "objects": objects} + path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + + +def _mesh_for_prim( + prim: ScenePrimMesh, + spec: CollisionSpec, +) -> o3d.geometry.TriangleMesh | None: + override = spec.resolve(prim.prim_path or prim.name) + override_type = override.get("type", spec.default) + if override_type == "skip": + return None + + mesh = _mesh_from_arrays( + prim.vertices.astype(np.float64), + prim.triangles.astype(np.int64), + ) + target_faces = int(override.get("target_faces") or 0) + if target_faces > 0 and len(mesh.triangles) > target_faces: + mesh = mesh.simplify_quadric_decimation(target_number_of_triangles=target_faces) + mesh.remove_degenerate_triangles() + mesh.remove_duplicated_triangles() + mesh.remove_duplicated_vertices() + mesh.remove_non_manifold_edges() + return mesh + + +def _mesh_from_arrays(vertices: np.ndarray, faces: np.ndarray) -> o3d.geometry.TriangleMesh: + mesh = o3d.geometry.TriangleMesh() + mesh.vertices = o3d.utility.Vector3dVector(vertices) + mesh.triangles = o3d.utility.Vector3iVector(faces.astype(np.int32)) + return mesh + + +__all__ = [ + "OBJECTS_SIDECAR_NAME", + "BrowserCollisionCookResult", + "cook_browser_collision", + "extract_scene_objects", +] diff --git a/dimos/experimental/pimsim/scene/cook.py b/dimos/experimental/pimsim/scene/cook.py new file mode 100644 index 0000000000..58eeacf898 --- /dev/null +++ b/dimos/experimental/pimsim/scene/cook.py @@ -0,0 +1,330 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Offline scene package cooker. + +This is intentionally not a DimOS runtime module. It prepares cooked scene +packages that runtime modules consume through normal config. +""" + +from __future__ import annotations + +import argparse +from dataclasses import asdict +import hashlib +import json +from pathlib import Path +from typing import Any + +from dimos.experimental.pimsim.scene.browser_collision import cook_browser_collision +from dimos.experimental.pimsim.scene.entity_collision import ( + COLLISION_DIR_NAME, + cook_entity_collision_hulls, +) +from dimos.experimental.pimsim.scene.inspect import inspect_scene_asset +from dimos.experimental.pimsim.scene.plan import build_scene_cook_plan +from dimos.experimental.pimsim.scene.sidecar import SceneCookSidecar +from dimos.experimental.pimsim.scene.visual_blender import cook_plan_visual_assets +from dimos.experimental.pimsim.scene.visual_glb import cook_browser_visual +from dimos.simulation.mujoco.collision_spec import CollisionSpec +from dimos.simulation.mujoco.scene_mesh_to_mjcf import load_or_bake +from dimos.simulation.scene_assets.mesh_scene import SceneMeshAlignment +from dimos.simulation.scene_assets.spec import ( + BrowserCollisionSpec, + BrowserVisualSpec, + MujocoSceneSpec, + SceneCookSpec, + ScenePackage, +) +from dimos.utils.data import get_data_dir +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +SCENE_PACKAGE_DIR = get_data_dir("scene_packages") +_PACKAGE_KEY_LEN = 12 +_COOK_VERSION = 4 + + +def cook_scene_package( + source_path: str | Path, + *, + output_dir: str | Path | None = None, + alignment: SceneMeshAlignment | None = None, + collision_spec: CollisionSpec | None = None, + cook_sidecar: SceneCookSidecar | None = None, + visual_spec: BrowserVisualSpec | None = None, + browser_collision_spec: BrowserCollisionSpec | None = None, + mujoco_spec: MujocoSceneSpec | None = None, + rebake: bool = False, +) -> ScenePackage: + """Cook one source scene into a robot-agnostic package. + + The package contains browser artifacts (visual + collision GLBs, + semantic ``objects.json``), per-entity GLBs, and a scene-only MuJoCo + wrapper. Robots are attached at runtime via ``MjSpec.attach()`` inside + ``MujocoSimModule.start``; the cooker never touches robot MJCFs. + """ + source = Path(source_path).expanduser().resolve() + if not source.exists(): + raise FileNotFoundError(f"scene source not found: {source}") + + align = alignment or SceneMeshAlignment() + visual = visual_spec or BrowserVisualSpec() + browser_collision = browser_collision_spec or BrowserCollisionSpec() + mujoco = mujoco_spec or MujocoSceneSpec() + cook_spec = SceneCookSpec( + source_path=source, + alignment=align, + browser_visual=visual, + browser_collision=browser_collision, + mujoco=mujoco, + ) + sidecar = cook_sidecar or SceneCookSidecar.auto_discover(source) + + package_dir = ( + Path(output_dir).expanduser().resolve() + if output_dir is not None + else SCENE_PACKAGE_DIR / _package_key(cook_spec, sidecar) + ) + browser_dir = package_dir / "browser" + mujoco_dir = package_dir / "mujoco" + package_dir.mkdir(parents=True, exist_ok=True) + + stats: dict[str, Any] = { + "source": inspect_scene_asset(source).to_json_dict(), + "cook_spec": _cook_spec_json(cook_spec), + "cook_version": _COOK_VERSION, + } + if sidecar.path is not None or sidecar.interactables: + stats["authored_sidecar"] = sidecar.to_json_dict() + + plan = build_scene_cook_plan( + source, + sidecar=sidecar, + alignment=align, + output_dir=package_dir, + collision_spec=collision_spec, + ) + stats["cook_plan"] = plan.to_json_dict() + + entities = plan.entities_metadata() + if entities: + stats["interactables"] = { + "count": len(entities), + "ids": [entity["id"] for entity in entities], + "static_visual_filter": "plan/blender", + } + + visual_source = source + # Only invoke Blender when at least one entity actually extracts from + # the source mesh; pure-synthetic sidecars (manip rigs) don't need it. + needs_blender = visual.enabled and any( + entity.visual_path is not None for entity in plan.entities + ) + if needs_blender: + visual_source = cook_plan_visual_assets( + source, + package_dir, + plan=plan, + rebake=rebake, + ) + + if mujoco.enabled: + hull_counts = _cook_entity_collision(entities, rebake=rebake) + if hull_counts: + stats["entity_collision"] = {"hulls_per_entity": hull_counts} + + visual_result = cook_browser_visual( + visual_source, + browser_dir, + spec=visual, + rebake=rebake, + ) + if visual_result is not None: + stats["browser_visual"] = { + "tool": visual_result.tool, + **visual_result.stats, + } + + browser_collision_result = cook_browser_collision( + source, + browser_dir, + alignment=SceneMeshAlignment(y_up=False), + spec=browser_collision, + collision_spec=plan.collision_spec, + rebake=rebake, + ) + if browser_collision_result is not None: + stats["browser_collision"] = browser_collision_result.stats + + mujoco_scene_path: Path | None = None + if mujoco.enabled: + mujoco_scene_path = load_or_bake( + scene_mesh_path=source, + alignment=align, + cache_root=mujoco_dir, + collision_spec=plan.collision_spec, + include_visual_mesh=mujoco.include_visual_mesh, + rebake=rebake, + ) + stats["mujoco"] = {"scene_path": str(mujoco_scene_path)} + + package = ScenePackage( + package_dir=package_dir, + source_path=source, + alignment=align, + visual_path=visual_result.path if visual_result else None, + browser_collision_path=browser_collision_result.path if browser_collision_result else None, + objects_path=browser_collision_result.objects_path if browser_collision_result else None, + mujoco_scene_path=mujoco_scene_path, + metadata_path=package_dir / "scene.meta.json", + entities=entities, + stats=stats, + ) + package.write_metadata() + logger.info("scene package cooked: %s", package.metadata_path) + return package + + +def _cook_entity_collision( + entities: list[dict[str, Any]], + *, + rebake: bool, +) -> dict[str, int]: + """Decompose every mesh entity's GLB into package collision hulls. + + Mutates the entity metadata in place, recording the hull files as + ``collision_paths`` so the runtime composer loads them from the + package instead of decomposing at boot. Returns hull counts by + entity id. + """ + hull_counts: dict[str, int] = {} + for entity in entities: + if entity.get("descriptor", {}).get("shape_hint") != "mesh": + continue + visual_path = entity.get("visual_path") + if not visual_path or not Path(visual_path).exists(): + logger.warning( + "entity %s: mesh entity has no cooked visual GLB; " + "no collision hulls (runtime falls back to AABB box)", + entity.get("id"), + ) + continue + hull_paths = cook_entity_collision_hulls( + visual_path, + Path(visual_path).parent / COLLISION_DIR_NAME, + rebake=rebake, + ) + if hull_paths: + entity["collision_paths"] = [str(path) for path in hull_paths] + hull_counts[str(entity.get("id"))] = len(hull_paths) + return hull_counts + + +def _package_key( + cook_spec: SceneCookSpec, + sidecar: SceneCookSidecar, +) -> str: + h = hashlib.sha256() + h.update(cook_spec.source_path.read_bytes()) + h.update(str(_COOK_VERSION).encode()) + h.update(json.dumps(_cook_spec_json(cook_spec), sort_keys=True).encode()) + h.update(json.dumps(sidecar.to_json_dict(), sort_keys=True).encode()) + return h.hexdigest()[:_PACKAGE_KEY_LEN] + + +def _cook_spec_json(cook_spec: SceneCookSpec) -> dict[str, Any]: + raw = asdict(cook_spec) + raw["source_path"] = str(cook_spec.source_path) + return raw + + +def _parse_xyz(value: str) -> tuple[float, float, float]: + parts = [float(part.strip()) for part in value.split(",")] + if len(parts) != 3: + raise argparse.ArgumentTypeError("expected comma-separated x,y,z") + return (parts[0], parts[1], parts[2]) + + +def cli_main() -> None: + parser = argparse.ArgumentParser( + description="Cook a scene asset into a robot-agnostic DimOS scene package.", + ) + parser.add_argument("source", type=Path) + parser.add_argument("--output-dir", type=Path) + parser.add_argument("--cook-spec", type=Path) + parser.add_argument("--scale", type=float, default=1.0) + parser.add_argument("--translation", type=_parse_xyz, default=(0.0, 0.0, 0.0)) + parser.add_argument("--rotation-zyx-deg", type=_parse_xyz, default=(0.0, 0.0, 0.0)) + parser.add_argument("--no-y-up", action="store_true") + parser.add_argument("--no-visual", action="store_true") + parser.add_argument( + "--visual-optimizer", + choices=("gltfpack", "blender", "copy"), + default="gltfpack", + ) + parser.add_argument("--visual-simplify-ratio", type=float, default=0.3) + parser.add_argument("--visual-simplify-error", type=float, default=0.02) + parser.add_argument("--visual-max-texture-size", type=int) + parser.add_argument( + "--visual-texture-format", + choices=("none", "webp", "ktx2"), + default="none", + ) + parser.add_argument("--no-browser-collision", action="store_true") + parser.add_argument("--browser-collision-target-faces", type=int, default=100_000) + parser.add_argument("--no-mujoco", action="store_true") + parser.add_argument("--include-mujoco-visual", action="store_true") + parser.add_argument("--rebake", action="store_true") + args = parser.parse_args() + + package = cook_scene_package( + args.source, + output_dir=args.output_dir, + alignment=SceneMeshAlignment( + scale=args.scale, + translation=args.translation, + rotation_zyx_deg=args.rotation_zyx_deg, + y_up=not args.no_y_up, + ), + cook_sidecar=SceneCookSidecar.from_json(args.cook_spec) if args.cook_spec else None, + visual_spec=BrowserVisualSpec( + enabled=not args.no_visual, + optimizer=args.visual_optimizer, + simplify_ratio=args.visual_simplify_ratio, + simplify_error=args.visual_simplify_error, + texture_format=( + None if args.visual_texture_format == "none" else args.visual_texture_format + ), + max_texture_size=args.visual_max_texture_size, + ), + browser_collision_spec=BrowserCollisionSpec( + enabled=not args.no_browser_collision, + target_faces=args.browser_collision_target_faces, + ), + mujoco_spec=MujocoSceneSpec( + enabled=not args.no_mujoco, + include_visual_mesh=args.include_mujoco_visual, + ), + rebake=args.rebake, + ) + print(package.metadata_path) + + +if __name__ == "__main__": + cli_main() + + +__all__ = ["SCENE_PACKAGE_DIR", "cook_scene_package"] diff --git a/dimos/experimental/pimsim/scene/entity_collision.py b/dimos/experimental/pimsim/scene/entity_collision.py new file mode 100644 index 0000000000..1ecd6619b3 --- /dev/null +++ b/dimos/experimental/pimsim/scene/entity_collision.py @@ -0,0 +1,141 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Cook-time convex decomposition for scene-package entities. + +Mesh entities need convex collision geometry for MuJoCo's narrowphase. +The cooker decomposes each entity's ``visual.glb`` with CoACD (chair +legs / seat / back each get their own hull, so contacts are +chair-shaped) and writes the hulls into the package next to the visual: + +```text +entities// +├── visual.glb +└── mujoco_collision/ + ├── hull_000.obj + └── ... +``` + +The hull paths are recorded per entity as ``collision_paths`` in +``scene.meta.json`` and loaded verbatim by the runtime composer +(``dimos/simulation/mujoco/entity_scene.py``) — there is no runtime +decomposition and no per-machine cache; the package is self-contained. +""" + +from __future__ import annotations + +from pathlib import Path + +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +COLLISION_DIR_NAME = "mujoco_collision" + +_COACD_MAX_HULLS = 32 +_COACD_THRESHOLD = 0.05 +_COACD_RESOLUTION = 500 +_COACD_MCTS_ITERATIONS = 30 +_COACD_MCTS_NODES = 10 + + +def cook_entity_collision_hulls( + visual_mesh_path: str | Path, + out_dir: str | Path, + *, + rebake: bool = False, +) -> list[Path]: + """Decompose one entity mesh into convex hulls under ``out_dir``. + + Idempotent: existing ``hull_*.obj`` files are reused unless ``rebake``. + Falls back to a single convex hull when CoACD is unavailable or fails + on the mesh. Returns ``[]`` (with a warning) when the mesh can't be + read at all — the runtime composer then uses an AABB box. + """ + mesh_path = Path(visual_mesh_path) + out_dir = Path(out_dir) + + if not rebake: + existing = sorted(out_dir.glob("hull_*.obj")) + if existing: + return existing + + try: + import open3d as o3d # type: ignore[import-untyped] + except ImportError as exc: + logger.warning("entity hulls: open3d unavailable (%s); skipping %s", exc, mesh_path) + return [] + + mesh = o3d.io.read_triangle_mesh(str(mesh_path)) + if not mesh.has_vertices(): + logger.warning("entity hulls: no vertices in %s; skipping", mesh_path) + return [] + + parts = _run_coacd(mesh, mesh_path) + + out_dir.mkdir(parents=True, exist_ok=True) + for stale in out_dir.glob("hull_*.obj"): + stale.unlink() + + import numpy as np + + out_paths: list[Path] = [] + if parts: + for i, (vertices, triangles) in enumerate(parts): + hull_mesh = o3d.geometry.TriangleMesh() + hull_mesh.vertices = o3d.utility.Vector3dVector(np.asarray(vertices, dtype=np.float64)) + hull_mesh.triangles = o3d.utility.Vector3iVector(np.asarray(triangles, dtype=np.int32)) + path = out_dir / f"hull_{i:03d}.obj" + o3d.io.write_triangle_mesh(str(path), hull_mesh, write_vertex_normals=False) + out_paths.append(path) + else: + hull, _ = mesh.compute_convex_hull() + path = out_dir / "hull_000.obj" + o3d.io.write_triangle_mesh(str(path), hull, write_vertex_normals=False) + out_paths.append(path) + + return out_paths + + +def _run_coacd(mesh: object, mesh_path: Path) -> list[tuple[object, object]]: + """CoACD parts for an open3d mesh; ``[]`` means fall back to one hull.""" + try: + import coacd # type: ignore[import-not-found, import-untyped] + import numpy as np + + if not getattr(_run_coacd, "_coacd_silenced", False): + coacd.set_log_level("error") + _run_coacd._coacd_silenced = True # type: ignore[attr-defined] + + cm = coacd.Mesh( + np.asarray(mesh.vertices, dtype=np.float64), # type: ignore[attr-defined] + np.asarray(mesh.triangles, dtype=np.int32), # type: ignore[attr-defined] + ) + parts = coacd.run_coacd( + cm, + threshold=_COACD_THRESHOLD, + max_convex_hull=_COACD_MAX_HULLS, + resolution=_COACD_RESOLUTION, + mcts_iterations=_COACD_MCTS_ITERATIONS, + mcts_nodes=_COACD_MCTS_NODES, + ) + return list(parts) + except Exception as exc: + logger.warning( + "entity hulls: CoACD failed for %s (%s); using single convex hull", mesh_path, exc + ) + return [] + + +__all__ = ["COLLISION_DIR_NAME", "cook_entity_collision_hulls"] diff --git a/dimos/experimental/pimsim/scene/inspect.py b/dimos/experimental/pimsim/scene/inspect.py new file mode 100644 index 0000000000..41158b548d --- /dev/null +++ b/dimos/experimental/pimsim/scene/inspect.py @@ -0,0 +1,185 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fast scene asset inspection for cook reports and budget checks.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import open3d as o3d # type: ignore[import-untyped] + + +@dataclass(frozen=True) +class SceneAssetStats: + path: str + bytes: int + format: str + mesh_count: int = 0 + node_count: int = 0 + material_count: int = 0 + texture_count: int = 0 + vertex_count: int = 0 + triangle_count: int = 0 + + def to_json_dict(self) -> dict[str, Any]: + return asdict(self) + + +def inspect_scene_asset(path: str | Path) -> SceneAssetStats: + """Return lightweight geometry/material counts for a supported scene file.""" + resolved = Path(path).expanduser().resolve() + if not resolved.exists(): + raise FileNotFoundError(f"scene asset not found: {resolved}") + + suffix = resolved.suffix.lower() + if suffix in {".glb", ".gltf"}: + return _inspect_gltf(resolved) + if suffix in {".usd", ".usda", ".usdc", ".usdz"}: + return _inspect_usd(resolved) + return _inspect_open3d(resolved) + + +def _inspect_gltf(path: Path) -> SceneAssetStats: + import trimesh + + loaded: Any = trimesh.load(str(path)) + if isinstance(loaded, trimesh.Trimesh): + # visual may be ColorVisuals (no material) or TextureVisuals. + material = getattr(loaded.visual, "material", None) + material_count = 1 if material is not None else 0 + return SceneAssetStats( + path=str(path), + bytes=path.stat().st_size, + format=path.suffix.lower().lstrip("."), + mesh_count=1, + node_count=1, + material_count=material_count, + texture_count=_count_material_textures([material]), + vertex_count=len(loaded.vertices), + triangle_count=len(loaded.faces), + ) + + scene = loaded + mesh_count = len(getattr(scene, "geometry", {})) + node_count = len(getattr(scene.graph, "nodes_geometry", [])) + materials = [] + vertex_count = 0 + triangle_count = 0 + for geom in scene.geometry.values(): + if not isinstance(geom, trimesh.Trimesh): + continue + vertex_count += len(geom.vertices) + triangle_count += len(geom.faces) + materials.append(getattr(geom.visual, "material", None)) + material_keys = {repr(material) for material in materials if material is not None} + return SceneAssetStats( + path=str(path), + bytes=path.stat().st_size, + format=path.suffix.lower().lstrip("."), + mesh_count=mesh_count, + node_count=node_count, + material_count=len(material_keys), + texture_count=_count_material_textures(materials), + vertex_count=vertex_count, + triangle_count=triangle_count, + ) + + +def _count_material_textures(materials: list[Any]) -> int: + textures: set[int] = set() + for material in materials: + if material is None: + continue + for name in ( + "baseColorTexture", + "metallicRoughnessTexture", + "normalTexture", + "emissiveTexture", + "occlusionTexture", + "image", + ): + image = getattr(material, name, None) + if image is not None: + textures.add(id(image)) + return len(textures) + + +def _inspect_usd(path: Path) -> SceneAssetStats: + try: + from pxr import Usd, UsdGeom, UsdShade # type: ignore[import-not-found, import-untyped] + except ImportError as exc: + raise ImportError("inspecting USD assets requires usd-core") from exc + + stage = Usd.Stage.Open(str(path)) + if stage is None: + raise RuntimeError(f"could not open USD stage: {path}") + + mesh_count = 0 + vertex_count = 0 + triangle_count = 0 + materials: set[str] = set() + textures: set[str] = set() + for prim in stage.Traverse(): + if prim.IsA(UsdGeom.Mesh): + mesh_count += 1 + mesh = UsdGeom.Mesh(prim) + points_raw = mesh.GetPointsAttr().Get() + counts_raw = mesh.GetFaceVertexCountsAttr().Get() + points = points_raw if points_raw is not None else [] + face_counts = np.asarray(counts_raw if counts_raw is not None else [], dtype=np.int32) + vertex_count += len(points) + triangle_count += int(np.maximum(face_counts - 2, 0).sum()) + bound = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()[0] + if bound: + materials.add(str(bound.GetPath())) + if prim.IsA(UsdShade.Shader): + shader = UsdShade.Shader(prim) + if shader.GetIdAttr().Get() == "UsdUVTexture": + file_input = shader.GetInput("file") + if file_input and file_input.Get() is not None: + textures.add(str(file_input.Get())) + + return SceneAssetStats( + path=str(path), + bytes=path.stat().st_size, + format=path.suffix.lower().lstrip("."), + mesh_count=mesh_count, + node_count=mesh_count, + material_count=len(materials), + texture_count=len(textures), + vertex_count=vertex_count, + triangle_count=triangle_count, + ) + + +def _inspect_open3d(path: Path) -> SceneAssetStats: + mesh = o3d.io.read_triangle_mesh(str(path)) + if len(mesh.triangles) == 0: + raise RuntimeError(f"empty mesh: {path}") + return SceneAssetStats( + path=str(path), + bytes=path.stat().st_size, + format=path.suffix.lower().lstrip("."), + mesh_count=1, + node_count=1, + vertex_count=len(mesh.vertices), + triangle_count=len(mesh.triangles), + ) + + +__all__ = ["SceneAssetStats", "inspect_scene_asset"] diff --git a/dimos/experimental/pimsim/scene/plan.py b/dimos/experimental/pimsim/scene/plan.py new file mode 100644 index 0000000000..68f81784e2 --- /dev/null +++ b/dimos/experimental/pimsim/scene/plan.py @@ -0,0 +1,341 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Resolved scene cook plan. + +The authored sidecar is user intent. The cook plan is resolved source-scene +membership: exactly which source prims become each runtime entity, which visual +nodes Blender must extract/delete, and which collision policy every downstream +cook must consume. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field, replace +from pathlib import Path +import re +from typing import Any + +import numpy as np + +from dimos.experimental.pimsim.scene.sidecar import InteractableSpec, SceneCookSidecar +from dimos.simulation.mujoco.collision_spec import CollisionSpec +from dimos.simulation.scene_assets.mesh_scene import ( + SceneMeshAlignment, + ScenePrimMesh, + load_scene_prims, +) + +_HASH_SUFFIX_RE = re.compile(r"_[0-9a-fA-F]{6,}$") + + +@dataclass(frozen=True) +class EntityCookPlan: + """Resolved authored entity.""" + + spec: InteractableSpec + safe_id: str + matched_prim_paths: tuple[str, ...] + visual_node_patterns: tuple[str, ...] + aabb_min: tuple[float, float, float] + aabb_max: tuple[float, float, float] + center: tuple[float, float, float] + initial_quat: tuple[float, float, float, float] + descriptor: dict[str, Any] + visual_path: Path | None + + def to_metadata(self) -> dict[str, Any]: + return { + "id": self.spec.id, + "tags": list(self.spec.tags), + "source_prim_paths": list(self.spec.source_prim_paths), + "matched_prim_paths": list(self.matched_prim_paths), + "visual_node_patterns": list(self.visual_node_patterns), + "remove_from_static": self.spec.remove_from_static, + "spawn": self.spec.spawn, + "synthetic": self.spec.is_synthetic, + "aabb": { + "min": list(self.aabb_min), + "max": list(self.aabb_max), + }, + "initial_pose": { + "x": self.center[0], + "y": self.center[1], + "z": self.center[2], + "qw": self.initial_quat[0], + "qx": self.initial_quat[1], + "qy": self.initial_quat[2], + "qz": self.initial_quat[3], + }, + "visual_path": str(self.visual_path) if self.visual_path else None, + "descriptor": self.descriptor, + "physics": self.spec.physics, + "visual": self.spec.visual, + } + + def to_json_dict(self) -> dict[str, Any]: + return { + "id": self.spec.id, + "safe_id": self.safe_id, + "matched_prim_paths": list(self.matched_prim_paths), + "visual_node_patterns": list(self.visual_node_patterns), + "aabb": {"min": list(self.aabb_min), "max": list(self.aabb_max)}, + "center": list(self.center), + "synthetic": self.spec.is_synthetic, + "descriptor": self.descriptor, + "visual_path": str(self.visual_path) if self.visual_path else None, + "remove_from_static": self.spec.remove_from_static, + } + + +@dataclass(frozen=True) +class SceneCookPlan: + """Resolved plan shared by every artifact writer.""" + + source_path: Path + alignment: SceneMeshAlignment + sidecar: SceneCookSidecar + collision_spec: CollisionSpec + entities: tuple[EntityCookPlan, ...] = () + stats: dict[str, Any] = field(default_factory=dict) + + @property + def has_entities(self) -> bool: + return bool(self.entities) + + def entities_metadata(self) -> list[dict[str, Any]]: + return [entity.to_metadata() for entity in self.entities] + + def to_json_dict(self) -> dict[str, Any]: + return { + "source_path": str(self.source_path), + "alignment": { + "scale": self.alignment.scale, + "rotation_zyx_deg": list(self.alignment.rotation_zyx_deg), + "translation": list(self.alignment.translation), + "y_up": self.alignment.y_up, + }, + "sidecar_path": str(self.sidecar.path) if self.sidecar.path else None, + "entities": [entity.to_json_dict() for entity in self.entities], + "stats": self.stats, + } + + +def build_scene_cook_plan( + source_path: str | Path, + *, + sidecar: SceneCookSidecar, + alignment: SceneMeshAlignment, + output_dir: str | Path, + collision_spec: CollisionSpec | None = None, +) -> SceneCookPlan: + source = Path(source_path).expanduser().resolve() + base_collision = collision_spec or sidecar.collision + if not sidecar.interactables: + return SceneCookPlan( + source_path=source, + alignment=alignment, + sidecar=sidecar, + collision_spec=base_collision, + stats={"source_prims": 0, "entities": 0}, + ) + + entities_dir = Path(output_dir).expanduser().resolve() / "entities" + needs_prims = any(item.source_prim_paths for item in sidecar.interactables) + prims = load_scene_prims(source, alignment=alignment) if needs_prims else [] + entities = tuple( + ( + _build_synthetic_entity_plan(item, entities_dir) + if item.is_synthetic + else _build_matched_entity_plan(item, prims, entities_dir) + ) + for item in sidecar.interactables + ) + effective_collision = _collision_spec_with_entity_skips(base_collision, entities) + return SceneCookPlan( + source_path=source, + alignment=alignment, + sidecar=sidecar, + collision_spec=effective_collision, + entities=entities, + stats={"source_prims": len(prims), "entities": len(entities)}, + ) + + +def _build_matched_entity_plan( + spec: InteractableSpec, + prims: list[ScenePrimMesh], + entities_dir: Path, +) -> EntityCookPlan: + matched = sorted( + (prim for prim in prims if spec.matches(prim)), + key=_prim_sort_key, + ) + if not matched: + patterns = ", ".join(spec.source_prim_paths) + raise ValueError(f"scene interactable {spec.id!r} matched no source prims: {patterns}") + + vertices = np.concatenate([prim.vertices for prim in matched], axis=0) + aabb_min_np = vertices.min(axis=0).astype(float) + aabb_max_np = vertices.max(axis=0).astype(float) + center_np = ((aabb_min_np + aabb_max_np) * 0.5).astype(float) + extents = np.maximum(aabb_max_np - aabb_min_np, 1e-4).astype(float) + safe_id = _safe_entity_id(spec.id) + visual_path = entities_dir / safe_id / "visual.glb" + + shape_hint, shape_extents = _resolve_shape(spec, extents) + descriptor = _make_descriptor(spec, shape_hint, shape_extents, visual_path) + + return EntityCookPlan( + spec=spec, + safe_id=safe_id, + matched_prim_paths=tuple(prim.prim_path or prim.name for prim in matched), + visual_node_patterns=_visual_node_patterns(matched), + aabb_min=(float(aabb_min_np[0]), float(aabb_min_np[1]), float(aabb_min_np[2])), + aabb_max=(float(aabb_max_np[0]), float(aabb_max_np[1]), float(aabb_max_np[2])), + center=(float(center_np[0]), float(center_np[1]), float(center_np[2])), + initial_quat=(1.0, 0.0, 0.0, 0.0), + descriptor=descriptor, + visual_path=visual_path, + ) + + +def _build_synthetic_entity_plan( + spec: InteractableSpec, + entities_dir: Path, +) -> EntityCookPlan: + """Synthetic entity: no source-prim extraction, primitive geometry, + pose from the spec. Used for manip rigs, test cubes, props you want + in the scene that aren't in the asset.""" + pose = spec.pose or {} + center = ( + float(pose.get("x", 0.0)), + float(pose.get("y", 0.0)), + float(pose.get("z", 0.0)), + ) + quat = ( + float(pose.get("qw", 1.0)), + float(pose.get("qx", 0.0)), + float(pose.get("qy", 0.0)), + float(pose.get("qz", 0.0)), + ) + extents_raw = spec.physics.get("extents") + if not extents_raw: + raise ValueError( + f"synthetic interactable {spec.id!r}: physics.extents required " + f"(no source mesh to derive bounds from)" + ) + extents_np = np.asarray([float(v) for v in extents_raw], dtype=float) + half = extents_np / 2.0 if len(extents_np) == 3 else extents_np + aabb_half = np.zeros(3, dtype=float) + aabb_half[: len(half)] = half[:3] if len(half) >= 3 else half + aabb_min = tuple(c - h for c, h in zip(center, aabb_half, strict=True)) + aabb_max = tuple(c + h for c, h in zip(center, aabb_half, strict=True)) + + shape_hint, shape_extents = _resolve_shape(spec, extents_np) + safe_id = _safe_entity_id(spec.id) + descriptor = _make_descriptor(spec, shape_hint, shape_extents, visual_path=None) + + return EntityCookPlan( + spec=spec, + safe_id=safe_id, + matched_prim_paths=(), + visual_node_patterns=(), + aabb_min=aabb_min, + aabb_max=aabb_max, + center=center, + initial_quat=quat, + descriptor=descriptor, + visual_path=None, + ) + + +def _resolve_shape( + spec: InteractableSpec, + extents_np: np.ndarray, +) -> tuple[str, list[float]]: + shape_hint = str(spec.physics.get("shape", "box")) + shape_extents = spec.physics.get("extents") + if shape_extents is not None: + return shape_hint, [float(v) for v in shape_extents] + if shape_hint == "box": + return shape_hint, [float(v) for v in extents_np[:3]] + if shape_hint == "sphere": + return shape_hint, [float(max(extents_np) * 0.5)] + if shape_hint == "cylinder": + return shape_hint, [ + float(max(extents_np[0], extents_np[1]) * 0.5), + float(extents_np[2] if len(extents_np) >= 3 else extents_np[-1]), + ] + return shape_hint, [] + + +def _make_descriptor( + spec: InteractableSpec, + shape_hint: str, + shape_extents: list[float], + visual_path: Path | None, +) -> dict[str, Any]: + descriptor: dict[str, Any] = { + "entity_id": spec.id, + "kind": spec.kind, + "shape_hint": shape_hint, + "extents": [float(value) for value in shape_extents], + "mass": float(spec.mass), + } + if visual_path is not None: + safe_id = _safe_entity_id(spec.id) + descriptor["mesh_ref"] = f"entities/{safe_id}/visual.glb" + else: + descriptor["mesh_ref"] = "" + rgba = spec.visual.get("rgba") if spec.visual else None + if isinstance(rgba, list | tuple) and len(rgba) == 4: + descriptor["rgba"] = [float(v) for v in rgba] + return descriptor + + +def _visual_node_patterns(prims: list[ScenePrimMesh]) -> tuple[str, ...]: + names: list[str] = [] + for prim in prims: + prim_path = prim.visual_node_name or prim.prim_path or prim.name + basename = prim_path.lstrip("/").rsplit("/", 1)[-1] + visual_name = _HASH_SUFFIX_RE.sub("", basename) + if visual_name not in names: + names.append(visual_name) + return tuple(names) + + +def _collision_spec_with_entity_skips( + collision_spec: CollisionSpec, + entities: tuple[EntityCookPlan, ...], +) -> CollisionSpec: + prim_overrides: dict[str, dict[str, Any]] = dict(collision_spec.prim_overrides) + for entity in entities: + if not entity.spec.remove_from_static: + continue + for prim_path in sorted(entity.matched_prim_paths): + prim_overrides.setdefault(prim_path, {"type": "skip"}) + return replace(collision_spec, prim_overrides=prim_overrides) + + +def _prim_sort_key(prim: ScenePrimMesh) -> tuple[str, str]: + return (prim.prim_path or prim.name, prim.name) + + +def _safe_entity_id(entity_id: str) -> str: + safe = "".join(c if c.isalnum() or c in {"-", "_", "."} else "_" for c in entity_id) + return safe or "entity" + + +__all__ = ["EntityCookPlan", "SceneCookPlan", "build_scene_cook_plan"] diff --git a/dimos/experimental/pimsim/scene/sidecar.py b/dimos/experimental/pimsim/scene/sidecar.py new file mode 100644 index 0000000000..492700199b --- /dev/null +++ b/dimos/experimental/pimsim/scene/sidecar.py @@ -0,0 +1,189 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Authored scene-cook sidecar. + +The existing ``.collision.json`` file remains the low-level collision +contract. ``.cook.json`` is the wider authored-scene contract: it can +carry the same collision policy plus a small, explicit list of objects that +should be removed from static cooks and respawned as pimsim entities. +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +import fnmatch +import json +from pathlib import Path +from typing import Any, Literal + +from dimos.simulation.mujoco.collision_spec import CollisionSpec +from dimos.simulation.scene_assets.mesh_scene import ScenePrimMesh +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +CookEntitySpawn = Literal["initial", "manual"] +CookEntityKind = Literal["dynamic", "kinematic", "static"] + +_COOK_SIDECAR_SUFFIXES = (".cook.json", ".scene.json") + + +@dataclass(frozen=True) +class InteractableSpec: + """One hand-authored runtime entity. + + Two flavours: + * ``source_prim_paths`` set -> matched against scene prims, geometry + extracted from the source mesh (the chairs flow). + * ``source_prim_paths`` empty + ``pose`` set -> synthetic entity + with no scene-mesh source. Geometry comes from ``physics.shape`` + + ``physics.extents``; pose comes from the spec directly. Used for + manip rigs, test cubes, anything not already in the source asset. + """ + + id: str + source_prim_paths: tuple[str, ...] = () + pose: dict[str, float] | None = None + remove_from_static: bool = True + spawn: CookEntitySpawn = "initial" + kind: CookEntityKind = "dynamic" + mass: float = 1.0 + tags: tuple[str, ...] = () + physics: dict[str, Any] = field(default_factory=dict) + visual: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_dict(cls, raw: dict[str, Any]) -> InteractableSpec: + prims = raw.get("source_prim_paths", raw.get("prim_paths", ())) + if isinstance(prims, str): + prims = (prims,) + pose = raw.get("pose") + if not prims and pose is None: + raise ValueError( + f"interactable {raw.get('id')!r}: needs either source_prim_paths " + f"(extract from scene) or pose (synthetic entity)" + ) + tags = raw.get("tags", ()) + if isinstance(tags, str): + tags = (tags,) + return cls( + id=str(raw["id"]), + source_prim_paths=tuple(str(pattern) for pattern in prims), + pose=dict(pose) if pose is not None else None, + remove_from_static=bool(raw.get("remove_from_static", True)), + spawn=raw.get("spawn", "initial"), + kind=raw.get("kind", "dynamic"), + mass=float(raw.get("mass", 1.0)), + tags=tuple(str(tag) for tag in tags), + physics=dict(raw.get("physics", {})), + visual=dict(raw.get("visual", {})), + ) + + @property + def is_synthetic(self) -> bool: + return not self.source_prim_paths and self.pose is not None + + def to_json_dict(self) -> dict[str, Any]: + raw = asdict(self) + raw["source_prim_paths"] = list(self.source_prim_paths) + raw["tags"] = list(self.tags) + return raw + + def matches(self, prim: ScenePrimMesh) -> bool: + prim_candidates = tuple( + candidate for candidate in (prim.visual_node_name, prim.prim_path) if candidate + ) + return any( + match_prim_pattern(candidate, pattern, include_sanitized=False) + for candidate in prim_candidates + for pattern in self.source_prim_paths + ) + + +@dataclass(frozen=True) +class SceneCookSidecar: + """Authored policy loaded from ``.cook.json``. + + ``collision`` is exactly the older collision sidecar schema. Interactables + add surgical scene knowledge without forcing every object in a scene to + become semantic metadata. + """ + + path: Path | None = None + collision: CollisionSpec = field(default_factory=CollisionSpec) + interactables: tuple[InteractableSpec, ...] = () + + @classmethod + def from_dict(cls, raw: dict[str, Any], *, path: Path | None = None) -> SceneCookSidecar: + collision_raw = raw.get("collision") + if isinstance(collision_raw, dict): + collision = CollisionSpec.from_dict(collision_raw) + else: + # Accept old collision keys at top-level so authored sidecars can be + # promoted without wrapping every existing key manually. + collision = CollisionSpec.from_dict(raw) + interactables = tuple( + InteractableSpec.from_dict(item) for item in raw.get("interactables", ()) + ) + return cls(path=path, collision=collision, interactables=interactables) + + @classmethod + def from_json(cls, path: str | Path) -> SceneCookSidecar: + sidecar_path = Path(path).expanduser().resolve() + return cls.from_dict(json.loads(sidecar_path.read_text()), path=sidecar_path) + + @classmethod + def auto_discover(cls, scene_path: str | Path) -> SceneCookSidecar: + source = Path(scene_path).expanduser().resolve() + for suffix in _COOK_SIDECAR_SUFFIXES: + sidecar = source.with_suffix(suffix) + if sidecar.exists(): + logger.info("loading scene cook sidecar: %s", sidecar) + return cls.from_json(sidecar) + + legacy_collision = source.with_suffix(".collision.json") + if legacy_collision.exists(): + logger.info("loading legacy collision sidecar: %s", legacy_collision) + return cls(path=legacy_collision, collision=CollisionSpec.from_json(legacy_collision)) + return cls() + + def to_json_dict(self) -> dict[str, Any]: + return { + "path": str(self.path) if self.path else None, + "collision": asdict(self.collision), + "interactables": [item.to_json_dict() for item in self.interactables], + } + + +def match_prim_pattern( + prim_path: str, + pattern: str, + *, + include_sanitized: bool = True, +) -> bool: + stripped = prim_path.lstrip("/") + sanitized = "".join(c if c.isalnum() else "_" for c in stripped) + basename = stripped.rsplit("/", 1)[-1] + candidates = [prim_path, stripped, basename] + if include_sanitized: + candidates.append(sanitized) + return any(fnmatch.fnmatchcase(candidate, pattern) for candidate in candidates) + + +__all__ = [ + "InteractableSpec", + "SceneCookSidecar", + "match_prim_pattern", +] diff --git a/dimos/experimental/pimsim/scene/test_entity_collision.py b/dimos/experimental/pimsim/scene/test_entity_collision.py new file mode 100644 index 0000000000..1e5ea81d05 --- /dev/null +++ b/dimos/experimental/pimsim/scene/test_entity_collision.py @@ -0,0 +1,67 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from dimos.experimental.pimsim.scene.entity_collision import cook_entity_collision_hulls + + +def _write_box_mesh(path: Path) -> None: + o3d = pytest.importorskip("open3d") + mesh = o3d.geometry.TriangleMesh.create_box(0.1, 0.1, 0.1) + o3d.io.write_triangle_mesh(str(path), mesh) + + +def test_cook_entity_collision_hulls(tmp_path: Path) -> None: + pytest.importorskip("coacd") + source = tmp_path / "visual.obj" + _write_box_mesh(source) + out_dir = tmp_path / "mujoco_collision" + + hulls = cook_entity_collision_hulls(source, out_dir) + + assert hulls, "expected at least one hull" + assert all(path.exists() for path in hulls) + assert all(path.parent == out_dir for path in hulls) + assert all(path.name.startswith("hull_") for path in hulls) + + +def test_cook_entity_collision_hulls_is_idempotent(tmp_path: Path) -> None: + pytest.importorskip("coacd") + source = tmp_path / "visual.obj" + _write_box_mesh(source) + out_dir = tmp_path / "mujoco_collision" + + first = cook_entity_collision_hulls(source, out_dir) + mtimes = [path.stat().st_mtime_ns for path in first] + + second = cook_entity_collision_hulls(source, out_dir) + assert second == first + assert [path.stat().st_mtime_ns for path in second] == mtimes + + rebaked = cook_entity_collision_hulls(source, out_dir, rebake=True) + assert rebaked + assert all(path.exists() for path in rebaked) + + +def test_cook_entity_collision_hulls_unreadable_mesh(tmp_path: Path) -> None: + pytest.importorskip("open3d") + source = tmp_path / "visual.obj" + source.write_text("not a mesh\n") + + assert cook_entity_collision_hulls(source, tmp_path / "mujoco_collision") == [] diff --git a/dimos/experimental/pimsim/scene/test_spec.py b/dimos/experimental/pimsim/scene/test_spec.py new file mode 100644 index 0000000000..24e7f67011 --- /dev/null +++ b/dimos/experimental/pimsim/scene/test_spec.py @@ -0,0 +1,324 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest + +from dimos.experimental.pimsim.scene import plan as plan_module +from dimos.experimental.pimsim.scene.sidecar import SceneCookSidecar +from dimos.simulation.scene_assets.mesh_scene import SceneMeshAlignment, ScenePrimMesh +from dimos.simulation.scene_assets.spec import ARTIFACT_FRAMES, ScenePackage, load_scene_package + + +def _metadata(tmp_path: Path) -> dict[str, object]: + return { + "source_path": str(tmp_path / "source.glb"), + "package_dir": str(tmp_path), + "alignment": { + "scale": 1.0, + "rotation_zyx_deg": [0.0, 0.0, 0.0], + "translation": [0.0, 0.0, 0.0], + "y_up": True, + }, + "artifact_frames": ARTIFACT_FRAMES, + "artifacts": { + "browser_visual": str(tmp_path / "visual.glb"), + "browser_collision": str(tmp_path / "collision.glb"), + "objects": str(tmp_path / "objects.json"), + "mujoco_scene": str(tmp_path / "wrapper.xml"), + }, + "stats": {}, + } + + +def test_load_scene_package_rejects_missing_artifact_frames(tmp_path: Path) -> None: + raw = _metadata(tmp_path) + raw.pop("artifact_frames") + metadata_path = tmp_path / "scene.meta.json" + metadata_path.write_text(json.dumps(raw)) + + with pytest.raises(ValueError, match="missing artifact frame metadata"): + load_scene_package(metadata_path) + + +def test_load_scene_package_rejects_mismatched_artifact_frames(tmp_path: Path) -> None: + raw = _metadata(tmp_path) + raw["artifact_frames"] = { + "browser_visual": "dimos_world", + "browser_collision": "source", + "mujoco": "dimos_world", + } + metadata_path = tmp_path / "scene.meta.json" + metadata_path.write_text(json.dumps(raw)) + + with pytest.raises(ValueError, match="artifact frame mismatch"): + load_scene_package(metadata_path) + + +def test_entity_collision_paths_round_trip(tmp_path: Path) -> None: + raw = _metadata(tmp_path) + raw["entities"] = [ + { + "id": "chair", + "visual_path": "entities/chair/visual.glb", + "collision_paths": [ + "entities/chair/mujoco_collision/hull_000.obj", + "entities/chair/mujoco_collision/hull_001.obj", + ], + "descriptor": {"shape_hint": "mesh"}, + } + ] + metadata_path = tmp_path / "scene.meta.json" + metadata_path.write_text(json.dumps(raw)) + + package = load_scene_package(metadata_path) + + assert package.entities[0]["collision_paths"] == [ + str(tmp_path / "entities/chair/mujoco_collision/hull_000.obj"), + str(tmp_path / "entities/chair/mujoco_collision/hull_001.obj"), + ] + assert package.to_json_dict()["entities"][0]["collision_paths"] == [ + "entities/chair/mujoco_collision/hull_000.obj", + "entities/chair/mujoco_collision/hull_001.obj", + ] + + +def test_load_scene_package_accepts_expected_artifact_frames(tmp_path: Path) -> None: + metadata_path = tmp_path / "scene.meta.json" + metadata_path.write_text(json.dumps(_metadata(tmp_path))) + + package = load_scene_package(metadata_path) + + assert package.visual_path == tmp_path / "visual.glb" + assert package.browser_collision_path == tmp_path / "collision.glb" + assert package.objects_path == tmp_path / "objects.json" + assert package.mujoco_scene_path == tmp_path / "wrapper.xml" + + +def test_scene_package_metadata_uses_package_relative_paths(tmp_path: Path) -> None: + package = ScenePackage( + package_dir=tmp_path, + source_path=tmp_path / "source.glb", + alignment=SceneMeshAlignment(), + visual_path=tmp_path / "browser" / "visual.glb", + browser_collision_path=tmp_path / "browser" / "collision.glb", + objects_path=tmp_path / "browser" / "objects.json", + mujoco_scene_path=tmp_path / "mujoco" / "abc123" / "wrapper.xml", + entities=[ + { + "id": "chair_001", + "visual_path": str(tmp_path / "entities" / "chair_001" / "visual.glb"), + } + ], + ) + + metadata_path = package.write_metadata() + raw = json.loads(metadata_path.read_text()) + + assert raw["package_dir"] == "." + assert raw["artifacts"]["browser_visual"] == "browser/visual.glb" + assert raw["artifacts"]["browser_collision"] == "browser/collision.glb" + assert raw["artifacts"]["objects"] == "browser/objects.json" + assert raw["artifacts"]["mujoco_scene"] == "mujoco/abc123/wrapper.xml" + assert raw["entities"][0]["visual_path"] == "entities/chair_001/visual.glb" + + loaded = load_scene_package(metadata_path) + assert loaded.package_dir == tmp_path + assert loaded.visual_path == tmp_path / "browser" / "visual.glb" + assert loaded.mujoco_scene_path == tmp_path / "mujoco" / "abc123" / "wrapper.xml" + assert loaded.entities[0]["visual_path"] == str( + tmp_path / "entities" / "chair_001" / "visual.glb" + ) + + +def test_load_scene_package_tolerates_missing_objects_sidecar(tmp_path: Path) -> None: + raw = _metadata(tmp_path) + # Older cooked packages without the semantic sidecar should still load. + raw["artifacts"].pop("objects") # type: ignore[union-attr] + metadata_path = tmp_path / "scene.meta.json" + metadata_path.write_text(json.dumps(raw)) + + package = load_scene_package(metadata_path) + + assert package.objects_path is None + + +def test_extract_scene_objects_emits_per_prim_aabb() -> None: + from dimos.experimental.pimsim.scene.browser_collision import extract_scene_objects + + triangles = np.array([[0, 1, 2]], dtype=np.int32) + prims = [ + ScenePrimMesh( + name="Sectional_seat", + prim_path="/Apt/Living/Sectional/seat", + vertices=np.array( + [[-1.0, -2.0, 0.0], [2.0, -2.0, 0.5], [-1.0, 1.0, 1.0]], + dtype=np.float32, + ), + triangles=triangles, + ), + ScenePrimMesh( + name="Empty_prim", + prim_path="/Apt/Living/Empty", + vertices=np.empty((0, 3), dtype=np.float32), + triangles=np.empty((0, 3), dtype=np.int32), + ), + ] + + objects = extract_scene_objects(prims) + + assert len(objects) == 1 # empty prim filtered + entry = objects[0] + assert entry["id"] == "Sectional_seat" + assert entry["prim_path"] == "/Apt/Living/Sectional/seat" + assert entry["aabb"]["min"] == [-1.0, -2.0, 0.0] + assert entry["aabb"]["max"] == [2.0, 1.0, 1.0] + + +def test_load_scene_package_preserves_packaged_entities(tmp_path: Path) -> None: + raw = _metadata(tmp_path) + raw["entities"] = [ + { + "id": "chair_016", + "descriptor": {"entity_id": "chair_016", "shape_hint": "box"}, + } + ] + metadata_path = tmp_path / "scene.meta.json" + metadata_path.write_text(json.dumps(raw)) + + package = load_scene_package(metadata_path) + + assert package.entities[0]["id"] == "chair_016" + + +def test_scene_cook_plan_maps_collision_prims_to_blender_visual_nodes( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + def fake_load_scene_prims( + path: str | Path, + alignment: SceneMeshAlignment | None = None, + ) -> list[ScenePrimMesh]: + del path, alignment + triangles = np.array([[0, 1, 2]], dtype=np.int32) + return [ + ScenePrimMesh( + name="Chair_seat", + prim_path="Chair_a1b2c3", + vertices=np.array( + [[-1.0, -1.0, 0.2], [-0.5, -1.0, 0.2], [-1.0, -0.5, 0.8]], + dtype=np.float32, + ), + triangles=triangles, + ), + ScenePrimMesh( + name="Chair.016_seat", + prim_path="Chair.016_a1b2c3", + vertices=np.array( + [[1.0, 2.0, 0.2], [2.0, 2.0, 0.2], [1.0, 3.0, 0.8]], + dtype=np.float32, + ), + triangles=triangles, + ), + ScenePrimMesh( + name="Chair.016_back", + prim_path="Chair.016_d4e5f6", + vertices=np.array( + [[1.0, 2.0, 0.8], [2.0, 3.0, 1.4], [1.5, 2.5, 1.2]], + dtype=np.float32, + ), + triangles=triangles, + ), + ] + + monkeypatch.setattr(plan_module, "load_scene_prims", fake_load_scene_prims) + sidecar = SceneCookSidecar.from_dict( + { + "interactables": [ + { + "id": "chair_000", + "source_prim_paths": ["Chair_*"], + "physics": {"shape": "box"}, + }, + { + "id": "chair_016", + "source_prim_paths": ["Chair.016_*"], + "physics": {"shape": "box"}, + }, + ] + } + ) + + plan = plan_module.build_scene_cook_plan( + tmp_path / "office.glb", + sidecar=sidecar, + alignment=SceneMeshAlignment(scale=2.0, y_up=False), + output_dir=tmp_path, + ) + + base_entity = plan.entities[0] + assert base_entity.matched_prim_paths == ("Chair_a1b2c3",) + assert base_entity.visual_node_patterns == ("Chair",) + assert base_entity.descriptor["mesh_ref"] == "entities/chair_000/visual.glb" + + entity = plan.entities[1] + assert entity.matched_prim_paths == ("Chair.016_a1b2c3", "Chair.016_d4e5f6") + assert entity.visual_node_patterns == ("Chair.016",) + assert entity.descriptor["mesh_ref"] == "entities/chair_016/visual.glb" + assert plan.collision_spec.resolve("Chair_a1b2c3")["type"] == "skip" + assert plan.collision_spec.resolve("Chair.016_a1b2c3")["type"] == "skip" + assert plan.collision_spec.resolve("Chair.001_a1b2c3")["type"] == "auto" + + +def test_synthetic_entity_uses_pose_and_extents(tmp_path: Path) -> None: + sidecar = SceneCookSidecar.from_dict( + { + "interactables": [ + { + "id": "manip_cube", + "pose": {"x": 0.0, "y": 0.75, "z": 0.69}, + "kind": "dynamic", + "mass": 0.15, + "physics": {"shape": "box", "extents": [0.08, 0.08, 0.08]}, + "visual": {"rgba": [0.85, 0.20, 0.20, 1.0]}, + "tags": ["manipulation"], + }, + ] + } + ) + plan = plan_module.build_scene_cook_plan( + tmp_path / "office.glb", + sidecar=sidecar, + alignment=SceneMeshAlignment(), + output_dir=tmp_path, + ) + + entity = plan.entities[0] + assert entity.spec.is_synthetic + assert entity.matched_prim_paths == () + assert entity.visual_path is None + assert entity.center == (0.0, 0.75, 0.69) + assert entity.descriptor["shape_hint"] == "box" + assert entity.descriptor["extents"] == [0.08, 0.08, 0.08] + assert entity.descriptor["rgba"] == [0.85, 0.20, 0.20, 1.0] + assert entity.descriptor["mesh_ref"] == "" + + +def test_interactable_requires_prims_or_pose() -> None: + with pytest.raises(ValueError, match="source_prim_paths.*or pose"): + SceneCookSidecar.from_dict({"interactables": [{"id": "ghost"}]}) diff --git a/dimos/experimental/pimsim/scene/visual_blender.py b/dimos/experimental/pimsim/scene/visual_blender.py new file mode 100644 index 0000000000..0615d271e1 --- /dev/null +++ b/dimos/experimental/pimsim/scene/visual_blender.py @@ -0,0 +1,313 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Blender-backed visual compiler for authored scene plans.""" + +from __future__ import annotations + +import json +from pathlib import Path +import shutil +import subprocess +import tempfile + +from dimos.experimental.pimsim.scene.plan import SceneCookPlan + +_COMMAND_TAIL_LINES = 30 +_VISUAL_PLAN_VERSION = 2 + +_PLAN_VISUAL_SCRIPT = r""" +import fnmatch +import json +import math +import pathlib +import sys + +import bpy +from mathutils import Matrix, Vector + +source = pathlib.Path(sys.argv[-2]) +plan_path = pathlib.Path(sys.argv[-1]) +plan = json.loads(plan_path.read_text()) +static_visual_path = pathlib.Path(plan["static_visual_path"]) +suffix = source.suffix.lower() + + +def fail(message): + raise RuntimeError(message) + + +def import_source(): + bpy.ops.object.select_all(action="SELECT") + bpy.ops.object.delete() + if suffix in {".usd", ".usda", ".usdc", ".usdz"}: + bpy.ops.wm.usd_import(filepath=str(source)) + elif suffix in {".gltf", ".glb"}: + bpy.ops.import_scene.gltf(filepath=str(source)) + elif suffix == ".obj": + bpy.ops.wm.obj_import(filepath=str(source)) + elif suffix == ".stl": + bpy.ops.wm.stl_import(filepath=str(source)) + elif suffix == ".ply": + bpy.ops.wm.ply_import(filepath=str(source)) + else: + fail(f"unsupported visual source suffix: {suffix}") + + +def alignment_matrix(): + align = plan["alignment"] + yaw, pitch, roll = [math.radians(float(v)) for v in align["rotation_zyx_deg"]] + cz, sz = math.cos(yaw), math.sin(yaw) + cy, sy = math.cos(pitch), math.sin(pitch) + cx, sx = math.cos(roll), math.sin(roll) + rz = Matrix(((cz, -sz, 0.0), (sz, cz, 0.0), (0.0, 0.0, 1.0))) + ry = Matrix(((cy, 0.0, sy), (0.0, 1.0, 0.0), (-sy, 0.0, cy))) + rx = Matrix(((1.0, 0.0, 0.0), (0.0, cx, -sx), (0.0, sx, cx))) + r = rz @ ry @ rx + if bool(align["y_up"]): + y_to_z = Matrix(((1.0, 0.0, 0.0), (0.0, 0.0, -1.0), (0.0, 1.0, 0.0))) + r = r @ y_to_z + return r, float(align["scale"]), Vector(tuple(float(v) for v in align["translation"])) + + +ALIGN_R, ALIGN_SCALE, ALIGN_T = alignment_matrix() + + +def object_candidates(obj): + candidates = {obj.name, obj.name_full} + if obj.parent is not None: + candidates.add(obj.parent.name) + candidates.add(obj.parent.name_full) + return candidates + + +def matches(obj, pattern): + return any(fnmatch.fnmatchcase(candidate, pattern) for candidate in object_candidates(obj)) + + +def source_point_from_blender_world(point): + if suffix in {".gltf", ".glb"}: + # Blender imports glTF Y-up coordinates into its Z-up world as + # (x, -z, y). The cook plan was resolved in the source glTF frame, + # so convert back before applying SceneMeshAlignment. + return Vector((point.x, point.z, -point.y)) + return point + + +def resolve_objects(entity): + objects = [] + missing = [] + for pattern in entity["visual_node_patterns"]: + matched = [obj for obj in bpy.context.scene.objects if obj.type == "MESH" and matches(obj, pattern)] + if not matched: + missing.append(pattern) + continue + for obj in matched: + if obj not in objects: + objects.append(obj) + if missing: + fail(f"entity {entity['id']} visual nodes not found in Blender import: {missing}") + return objects + + +def aligned_local_point(source_world, center): + source_point = source_point_from_blender_world(source_world) + return (ALIGN_R @ (ALIGN_SCALE * source_point)) + ALIGN_T - center + + +def duplicate_for_entity(obj, center, suffix): + dup = obj.copy() + dup.data = obj.data.copy() + dup.animation_data_clear() + dup.name = f"{obj.name}__{suffix}" + bpy.context.collection.objects.link(dup) + for vertex in dup.data.vertices: + source_world = obj.matrix_world @ vertex.co + vertex.co = aligned_local_point(source_world, center) + dup.matrix_world = Matrix.Identity(4) + return dup + + +def export_entity(entity, objects): + visual_path = pathlib.Path(entity["visual_path"]) + visual_path.parent.mkdir(parents=True, exist_ok=True) + center = Vector(tuple(float(v) for v in entity["center"])) + duplicates = [duplicate_for_entity(obj, center, entity["safe_id"]) for obj in objects] + try: + bpy.ops.object.select_all(action="DESELECT") + for dup in duplicates: + dup.select_set(True) + bpy.context.view_layer.objects.active = duplicates[0] + bpy.ops.export_scene.gltf( + filepath=str(visual_path), + export_format="GLB", + use_selection=True, + export_yup=False, + export_apply=True, + ) + finally: + for dup in duplicates: + bpy.data.objects.remove(dup, do_unlink=True) + + +def export_static_visual(objects_to_remove): + for obj in sorted(objects_to_remove, key=lambda item: item.name): + if obj.name in bpy.data.objects: + bpy.data.objects.remove(obj, do_unlink=True) + remaining = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"] + if not remaining: + fail("static visual would contain no mesh objects after entity removal") + static_visual_path.parent.mkdir(parents=True, exist_ok=True) + bpy.ops.object.select_all(action="DESELECT") + bpy.ops.export_scene.gltf( + filepath=str(static_visual_path), + export_format="GLB", + export_yup=True, + export_apply=True, + ) + + +import_source() +remove_from_static = set() +report = [] +for entity in plan["entities"]: + objects = resolve_objects(entity) + export_entity(entity, objects) + if entity["remove_from_static"]: + remove_from_static.update(objects) + report.append( + { + "id": entity["id"], + "visual_nodes": [obj.name for obj in objects], + "removed_from_static": entity["remove_from_static"], + } + ) +export_static_visual(remove_from_static) +print("DIMOS_VISUAL_PLAN_REPORT=" + json.dumps(report, sort_keys=True)) +""" + + +def cook_plan_visual_assets( + source_path: str | Path, + output_dir: str | Path, + *, + plan: SceneCookPlan, + rebake: bool = False, +) -> Path: + """Extract entity visuals and write a filtered static visual source GLB.""" + source = Path(source_path).expanduser().resolve() + package_dir = Path(output_dir).expanduser().resolve() + static_visual_source = package_dir / "browser" / "static_visual_source.glb" + plan_manifest = package_dir / "browser" / "visual_plan.json" + plan_json = _blender_plan_json(plan, static_visual_source) + expected_paths = [ + static_visual_source, + *(entity.visual_path for entity in plan.entities if entity.visual_path is not None), + ] + if ( + expected_paths + and all(path.exists() for path in expected_paths) + and _manifest_matches(plan_manifest, plan_json) + and not rebake + ): + return static_visual_source + + blender = shutil.which("blender") + if blender is None: + raise RuntimeError("authored visual entity extraction requires Blender on PATH") + + with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as plan_file: + json.dump(plan_json, plan_file) + plan_path = Path(plan_file.name) + with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as script_file: + script_file.write(_PLAN_VISUAL_SCRIPT) + script_path = Path(script_file.name) + try: + _run_command( + [ + blender, + "--background", + "--factory-startup", + "--python", + str(script_path), + "--", + str(source), + str(plan_path), + ], + "blender visual plan cook", + ) + finally: + plan_path.unlink(missing_ok=True) + script_path.unlink(missing_ok=True) + + plan_manifest.parent.mkdir(parents=True, exist_ok=True) + plan_manifest.write_text(json.dumps(plan_json, indent=2, sort_keys=True) + "\n") + return static_visual_source + + +def _blender_plan_json(plan: SceneCookPlan, static_visual_source: Path) -> dict[str, object]: + return { + "visual_plan_version": _VISUAL_PLAN_VERSION, + "alignment": { + "scale": plan.alignment.scale, + "rotation_zyx_deg": list(plan.alignment.rotation_zyx_deg), + "translation": list(plan.alignment.translation), + "y_up": plan.alignment.y_up, + }, + "static_visual_path": str(static_visual_source), + "entities": [ + { + "id": entity.spec.id, + "safe_id": entity.safe_id, + "visual_node_patterns": list(entity.visual_node_patterns), + "center": list(entity.center), + "visual_path": str(entity.visual_path), + "remove_from_static": entity.spec.remove_from_static, + } + # Synthetic entities have no source mesh to extract; skip them + # entirely so Blender doesn't try to match an empty pattern set. + for entity in plan.entities + if entity.visual_path is not None + ], + } + + +def _run_command(args: list[str], label: str) -> str: + result = subprocess.run( + args, + check=False, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + output = result.stdout or "" + if result.returncode != 0: + raise RuntimeError(f"{label} failed with exit code {result.returncode}:\n{_tail(output)}") + return output + + +def _manifest_matches(path: Path, expected: dict[str, object]) -> bool: + if not path.exists(): + return False + try: + return bool(json.loads(path.read_text()) == expected) + except json.JSONDecodeError: + return False + + +def _tail(output: str) -> str: + return "\n".join(output.splitlines()[-_COMMAND_TAIL_LINES:]) + + +__all__ = ["cook_plan_visual_assets"] diff --git a/dimos/experimental/pimsim/scene/visual_glb.py b/dimos/experimental/pimsim/scene/visual_glb.py new file mode 100644 index 0000000000..4a2027d10c --- /dev/null +++ b/dimos/experimental/pimsim/scene/visual_glb.py @@ -0,0 +1,363 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Cook browser visual assets for real-time browser rendering.""" + +from __future__ import annotations + +from dataclasses import dataclass +import json +from pathlib import Path +import shutil +import subprocess +import tempfile +from typing import Any + +from dimos.experimental.pimsim.scene.inspect import inspect_scene_asset +from dimos.simulation.scene_assets.spec import BrowserVisualSpec +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +_BLENDER_INPUT_SUFFIXES = { + ".usd", + ".usda", + ".usdc", + ".usdz", + ".gltf", + ".glb", + ".obj", + ".stl", + ".ply", +} +_GLTFPACK_INPUT_SUFFIXES = {".gltf", ".glb", ".obj"} +_COMMAND_TAIL_LINES = 30 + +_BLENDER_SCRIPT = r""" +import pathlib +import sys + +import bpy + +source = pathlib.Path(sys.argv[-4]) +target = pathlib.Path(sys.argv[-3]) +simplify_ratio = float(sys.argv[-2]) +max_texture_size = int(sys.argv[-1]) +suffix = source.suffix.lower() + +bpy.ops.object.select_all(action="SELECT") +bpy.ops.object.delete() + +if suffix in {".usd", ".usda", ".usdc", ".usdz"}: + bpy.ops.wm.usd_import(filepath=str(source)) +elif suffix in {".gltf", ".glb"}: + bpy.ops.import_scene.gltf(filepath=str(source)) +elif suffix == ".obj": + bpy.ops.wm.obj_import(filepath=str(source)) +elif suffix == ".stl": + bpy.ops.wm.stl_import(filepath=str(source)) +elif suffix == ".ply": + bpy.ops.wm.ply_import(filepath=str(source)) +else: + raise RuntimeError(f"unsupported visual source suffix: {suffix}") + +for obj in list(bpy.context.scene.objects): + if obj.type != "MESH": + bpy.data.objects.remove(obj, do_unlink=True) + +if max_texture_size > 0: + for image in bpy.data.images: + width, height = image.size + largest = max(width, height) + if largest <= max_texture_size: + continue + scale = max_texture_size / largest + try: + image.scale(max(1, int(width * scale)), max(1, int(height * scale))) + except RuntimeError: + # Blender cannot scale some generated or missing images; keep those + # untouched instead of aborting the entire scene cook. + pass + +if 0.0 < simplify_ratio < 0.999: + for obj in list(bpy.context.scene.objects): + if obj.type != "MESH": + continue + bpy.ops.object.select_all(action="DESELECT") + obj.select_set(True) + bpy.context.view_layer.objects.active = obj + modifier = obj.modifiers.new("dimos_decimate", "DECIMATE") + modifier.ratio = simplify_ratio + try: + bpy.ops.object.modifier_apply(modifier=modifier.name) + except RuntimeError: + obj.modifiers.remove(modifier) + +mesh_objects = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"] +if len(mesh_objects) > 1: + bpy.ops.object.select_all(action="DESELECT") + for obj in mesh_objects: + obj.select_set(True) + bpy.context.view_layer.objects.active = mesh_objects[0] + bpy.ops.object.join() + +bpy.ops.export_scene.gltf( + filepath=str(target), + export_format="GLB", + export_yup=True, + export_apply=True, +) +""" + + +@dataclass(frozen=True) +class BrowserVisualCookResult: + path: Path + stats: dict[str, Any] + tool: str + + +def cook_browser_visual( + source_path: str | Path, + output_dir: str | Path, + *, + spec: BrowserVisualSpec | None = None, + rebake: bool = False, +) -> BrowserVisualCookResult | None: + """Write the browser visual GLB for a scene package. + + ``gltfpack`` is the default path because browser performance is dominated + by draw calls, scene nodes, decoded texture memory, and shader/material + switches. Blender is kept as a conversion fallback for formats gltfpack + does not read directly. + """ + visual_spec = spec or BrowserVisualSpec() + if not visual_spec.enabled: + return None + + source = Path(source_path).expanduser().resolve() + out_dir = Path(output_dir).expanduser().resolve() + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / visual_spec.output_name + if out_path.exists() and not rebake: + return BrowserVisualCookResult( + path=out_path, + stats=inspect_scene_asset(out_path).to_json_dict(), + tool="cache", + ) + + source_stats = inspect_scene_asset(source).to_json_dict() + with tempfile.TemporaryDirectory(prefix="dimos-visual-cook-") as temp_dir_raw: + temp_dir = Path(temp_dir_raw) + temp_out = temp_dir / out_path.name + tool, report = _cook_visual(source, temp_out, visual_spec) + stats = inspect_scene_asset(temp_out).to_json_dict() + _validate_output(source_stats, stats, visual_spec) + if report is not None: + stats["optimizer_report"] = report + shutil.move(str(temp_out), out_path) + stats["path"] = str(out_path) + + warnings = _budget_warnings(stats, visual_spec) + if warnings: + stats["warnings"] = warnings + for warning in warnings: + logger.warning("browser visual budget: %s", warning) + return BrowserVisualCookResult(path=out_path, stats=stats, tool=tool) + + +def _cook_visual( + source: Path, + target: Path, + spec: BrowserVisualSpec, +) -> tuple[str, dict[str, Any] | None]: + optimizer = spec.optimizer.lower() + if optimizer == "copy": + if source.suffix.lower() != ".glb": + raise RuntimeError("copy visual optimizer requires a GLB source") + shutil.copy2(source, target) + return ("copy", None) + if optimizer == "blender": + _export_with_blender( + source, + target, + simplify_ratio=spec.simplify_ratio, + max_texture_size=spec.max_texture_size, + ) + return ("blender", None) + if optimizer == "gltfpack": + return _export_with_gltfpack(source, target, spec) + raise ValueError(f"unknown browser visual optimizer: {spec.optimizer}") + + +def _export_with_blender( + source: Path, + target: Path, + *, + simplify_ratio: float = 1.0, + max_texture_size: int | None = None, +) -> None: + blender = shutil.which("blender") + if blender is None: + raise RuntimeError( + f"{source.suffix} visual export requires Blender on PATH. Install Blender." + ) + + with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as script: + script.write(_BLENDER_SCRIPT) + script_path = Path(script.name) + try: + _run_command( + [ + blender, + "--background", + "--factory-startup", + "--python", + str(script_path), + "--", + str(source), + str(target), + str(simplify_ratio), + str(max_texture_size or 0), + ], + "blender", + ) + finally: + script_path.unlink(missing_ok=True) + + +def _export_with_gltfpack( + source: Path, + target: Path, + spec: BrowserVisualSpec, +) -> tuple[str, dict[str, Any] | None]: + command = _gltfpack_command() + source_for_gltfpack = source + with tempfile.TemporaryDirectory(prefix="dimos-gltfpack-source-") as temp_dir_raw: + if source.suffix.lower() not in _GLTFPACK_INPUT_SUFFIXES: + if source.suffix.lower() not in _BLENDER_INPUT_SUFFIXES: + raise RuntimeError(f"unsupported visual source suffix: {source.suffix}") + source_for_gltfpack = Path(temp_dir_raw) / "source.glb" + _export_with_blender(source, source_for_gltfpack) + + report_path = target.with_suffix(".gltfpack.json") + args = [ + *command, + "-i", + str(source_for_gltfpack), + "-o", + str(target), + "-mm", + "-si", + str(spec.simplify_ratio), + "-se", + str(spec.simplify_error), + "-r", + str(report_path), + ] + if spec.texture_format == "webp": + args.append("-tw") + elif spec.texture_format == "ktx2": + args.append("-tc") + elif spec.texture_format is not None: + raise ValueError(f"unknown browser texture format: {spec.texture_format}") + if spec.max_texture_size is not None: + if spec.texture_format is None: + raise ValueError("max_texture_size requires texture_format='webp' or 'ktx2'") + args.extend(["-tl", str(spec.max_texture_size)]) + + output = _run_command(args, "gltfpack") + if output and "Warning:" in output: + logger.warning("gltfpack output:\n%s", _tail(output)) + report = _read_json(report_path) + return ("gltfpack", report) + + +def _gltfpack_command() -> list[str]: + gltfpack = shutil.which("gltfpack") + if gltfpack is not None: + return [gltfpack] + npx = shutil.which("npx") + if npx is not None: + return [npx, "-y", "gltfpack"] + raise RuntimeError( + "browser visual optimization requires gltfpack. Install it with " + "`npm install -g gltfpack` or use --visual-optimizer blender/copy." + ) + + +def _run_command(args: list[str], label: str) -> str: + result = subprocess.run( + args, + check=False, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + output = result.stdout or "" + if result.returncode != 0: + raise RuntimeError(f"{label} failed with exit code {result.returncode}:\n{_tail(output)}") + return output + + +def _read_json(path: Path) -> dict[str, Any] | None: + if not path.exists(): + return None + try: + data = json.loads(path.read_text()) + except json.JSONDecodeError: + logger.warning("failed to parse optimizer report: %s", path) + return None + return data if isinstance(data, dict) else None + + +def _validate_output( + source_stats: dict[str, Any], + output_stats: dict[str, Any], + spec: BrowserVisualSpec, +) -> None: + source_vertices = int(source_stats.get("vertex_count") or 0) + output_vertices = int(output_stats.get("vertex_count") or 0) + if source_vertices <= 0 or output_vertices <= 0: + return + max_vertices = int(source_vertices * spec.max_vertex_growth_ratio) + if output_vertices > max_vertices: + raise RuntimeError( + "browser visual cook increased vertex count from " + f"{source_vertices} to {output_vertices}; refusing to write worse asset" + ) + + +def _tail(output: str) -> str: + return "\n".join(output.splitlines()[-_COMMAND_TAIL_LINES:]) + + +def _budget_warnings(stats: dict[str, Any], spec: BrowserVisualSpec) -> list[str]: + warnings: list[str] = [] + mesh_count = int(stats.get("node_count") or stats.get("mesh_count") or 0) + material_count = int(stats.get("material_count") or 0) + texture_count = int(stats.get("texture_count") or 0) + vertex_count = int(stats.get("vertex_count") or 0) + if mesh_count > spec.max_meshes: + warnings.append(f"{mesh_count} render nodes exceeds target {spec.max_meshes}") + if material_count > spec.max_materials: + warnings.append(f"{material_count} materials exceeds target {spec.max_materials}") + if texture_count > spec.max_textures: + warnings.append(f"{texture_count} textures exceeds target {spec.max_textures}") + if vertex_count > spec.max_vertices: + warnings.append(f"{vertex_count} vertices exceeds target {spec.max_vertices}") + return warnings + + +__all__ = ["BrowserVisualCookResult", "cook_browser_visual"] diff --git a/dimos/simulation/mujoco/collision_spec.py b/dimos/simulation/mujoco/collision_spec.py new file mode 100644 index 0000000000..a11427352b --- /dev/null +++ b/dimos/simulation/mujoco/collision_spec.py @@ -0,0 +1,968 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-prim collision-shape decision-making for ``bake_scene_mjcf``. + +The bake's job is to turn each USD/glTF/OBJ prim into one or more MuJoCo +````s. This module separates the *decision* (what shape to emit) +from the *emission* (the OBJ/MJCF writing). Three layers cooperate: + +1. **Generic geometric heuristics** — applied to every prim regardless of + source. Tiny-prim skip, aspect-ratio box override, near-convex check. + Safe defaults; no scene-specific knowledge. + +2. **Primitive auto-fit** — try OBB box / Ritter sphere / PCA cylinder / + PCA capsule. Accept the best fit if + ``hull_volume / primitive_volume >= fill_threshold``. Geometric only. + +3. **Sidecar overrides** — a JSON file (``.collision.json`` next + to the source mesh, or explicit path) with ``fnmatch`` patterns over + USD prim paths. Lets users skip lamps, force cylinders on pillars, + tune CoACD per pattern. Whoever produces this file (a human, a + future UE-side extractor, an LLM…) doesn't matter to the bake — the + sidecar is the contract. + +The dispatcher ``emit_for_prim()`` walks: sidecar override → generic +heuristics → primitive auto-fit → CoACD fallback, and returns a list +of ``GeomEmission`` describing every ```` the wrapper should +include for the prim. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import fnmatch +import json +from pathlib import Path +from typing import Any, Literal + +import numpy as np +from scipy.spatial import ConvexHull, QhullError # type: ignore[import-untyped] + +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +PrimitiveFit = dict[str, Any] +OverrideConfig = dict[str, Any] + + +# --------------------------------------------------------------------------- # +# Sidecar spec dataclasses # +# --------------------------------------------------------------------------- # + + +@dataclass +class CollisionSpec: + """User-facing collision configuration loaded from ``.collision.json``. + + Patterns in ``prim_overrides`` are matched with ``fnmatch`` (Unix-shell + globs) against the USD prim path of each prim — e.g. ``/Root/SM_Pillar_*``. + First-match wins; iteration order of the dict is preserved (Python 3.7+). + + Each override value is a dict with at minimum ``"type"``: + + - ``"box"`` / ``"sphere"`` / ``"cylinder"`` / ``"capsule"`` / ``"plane"``: + force the corresponding primitive. Auto-fit picks the parameters + unless explicit ``"size"`` / ``"pos"`` / ``"quat"`` is provided. + - ``"hull"``: force single convex hull, no CoACD. + - ``"mesh"`` / ``"decimate"``: emit this prim as a mesh geom. Optional + ``"target_faces"`` simplifies the source mesh before MJCF emission. + - ``"decompose"``: force CoACD even if auto-fit would have accepted a + primitive. Optional ``"max_hulls"`` overrides ``coacd_max_hulls``. + - ``"skip"``: emit no collision geom. Visual mesh still drawn. + - ``"auto"``: same as the global default (useful to scope a pattern + back to default behaviour inside a wider override). + + Optional override keys: + + - ``"friction"``: list ``[slide, spin, roll]``. + - ``"min_thickness"``: for ``"box"`` overrides, minimum full + thickness in metres along world Z. Useful for authored walkable + floor slabs exported as zero-thickness sheets. + - ``"preserve"``: with ``"min_thickness"``, which world-Z face to + keep fixed while expanding: ``"top"``, ``"bottom"``, or + ``"center"``. Defaults to ``"center"``. + - ``"max_hulls"``: per-pattern CoACD cap. + - ``"target_faces"``: per-pattern triangle target for ``mesh`` / + ``decimate`` outputs, or a post-process cap for hull outputs. + """ + + #: Fallback policy when no pattern matches. ``"auto"`` runs the full + #: heuristics→primitive→CoACD pipeline. ``"hull"`` always emits one + #: convex hull. ``"skip"`` emits nothing (visual only). + default: Literal["auto", "hull", "skip"] = "auto" + + #: A primitive is accepted in auto-fit if + #: ``hull_volume / primitive_volume >= fill_threshold``. Higher = + #: stricter (more prims fall through to CoACD). + fill_threshold: float = 0.85 + + #: Prims whose largest extent is below this (metres) emit no geom. + #: Catches trim/fasteners that the robot can't meaningfully contact. + tiny_prim_extent_m: float = 0.03 + + #: If one axis is ``>= aspect_ratio_box`` times the smaller two, the + #: prim is forced to a box even if auto-fit fill ratio is borderline. + #: Catches wall panels, floor slabs, doors. + aspect_ratio_box: float = 20.0 + + #: If mesh's hull is this close to its actual mesh volume, use one + #: hull and skip CoACD entirely (mesh is already near-convex). + near_convex_threshold: float = 0.9 + + #: CoACD concavity threshold (URLab default). Lower = finer + #: decomposition (more sub-hulls). + coacd_threshold: float = 0.05 + + #: Hard cap on hulls per CoACD invocation. -1 = unlimited. + coacd_max_hulls: int = 64 + + #: Only run decomposition when the prim's single-hull volume exceeds + #: this (m³). Smaller furniture-scale prims use one hull regardless. + shell_volume_m3: float = 2.0 + + #: Preserve large non-rectangular sheet footprints with thin triangle + #: prisms. This helps moderate indoor scenes with angular floors, but + #: is disabled by ``bake_scene_mjcf`` for very large scenes unless + #: explicitly overridden. + enable_sheet_prisms: bool = True + + #: Scene-level guard used by ``bake_scene_mjcf``. Above this many + #: source prims, sheet prisms can explode the geom count; use sidecar + #: overrides for specific floors instead. + sheet_prism_max_scene_prims: int = 2500 + + #: ``USD-path-glob -> override-dict``. See class docstring. + prim_overrides: dict[str, OverrideConfig] = field(default_factory=dict) + + #: Split suspicious scene-graph nodes that are really many disconnected + #: tiny meshes spread over a large area before running primitive fitting. + split_disconnected_components: bool = True + + #: Minimum component count before a prim is considered grouped clutter. + split_min_components: int = 8 + + #: Combined prim extent must be this much larger than the median component + #: extent before splitting. This avoids splitting normal multi-part props. + split_extent_ratio: float = 4.0 + + #: A prim must span at least this far before the cooker spends time + #: checking disconnected components. This keeps normal props cheap. + split_prim_min_extent_m: float = 5.0 + + #: The prim must also be slab-like by axis ratio before splitting. + #: This targets the path that can otherwise emit one giant box. + split_axis_ratio: float = 20.0 + + #: Components smaller than this largest extent are dropped after splitting. + #: This catches leaves, cups, bottles, and other decorative clutter that can + #: destabilize navigation collision while being too small to matter. + split_component_min_extent_m: float = 0.15 + + #: Very small triangle islands are dropped after splitting. + split_component_min_faces: int = 16 + + @classmethod + def from_dict(cls, raw: dict[str, Any]) -> CollisionSpec: + """Build a collision spec from decoded JSON. + + Unknown keys are ignored to keep authored sidecars forwards-compatible. + """ + known = { + "default", + "fill_threshold", + "tiny_prim_extent_m", + "aspect_ratio_box", + "near_convex_threshold", + "coacd_threshold", + "coacd_max_hulls", + "shell_volume_m3", + "enable_sheet_prisms", + "sheet_prism_max_scene_prims", + "prim_overrides", + "split_disconnected_components", + "split_min_components", + "split_extent_ratio", + "split_prim_min_extent_m", + "split_axis_ratio", + "split_component_min_extent_m", + "split_component_min_faces", + } + kwargs = {k: v for k, v in raw.items() if k in known} + # Ignore "$schema" and any future top-level keys silently. + return cls(**kwargs) + + @classmethod + def from_json(cls, path: Path | str) -> CollisionSpec: + """Load a sidecar. Unknown keys are ignored to keep the format forwards-compatible.""" + path = Path(path) + raw = json.loads(path.read_text()) + return cls.from_dict(raw) + + @classmethod + def auto_discover(cls, scene_path: Path | str) -> CollisionSpec: + """Return the sidecar next to ``scene_path`` if it exists, else defaults.""" + scene_path = Path(scene_path) + sidecar = scene_path.with_suffix(".collision.json") + if sidecar.exists(): + logger.info(f"loading collision sidecar: {sidecar}") + return cls.from_json(sidecar) + return cls() + + def resolve(self, prim_path: str) -> OverrideConfig: + """Find the matching override for ``prim_path`` (USD path). + + Returns a dict with at least ``"type"``. Falls back to + ``{"type": self.default}`` when no pattern matches. + """ + stripped = prim_path.lstrip("/") + sanitized = "".join(c if c.isalnum() else "_" for c in stripped) + candidates = (prim_path, stripped, sanitized) + for pattern, override in self.prim_overrides.items(): + if any(fnmatch.fnmatchcase(candidate, pattern) for candidate in candidates): + # Pattern's "auto" defers to global default. + if override.get("type") == "auto": + return {**override, "type": self.default} + return override + return {"type": self.default} + + +# --------------------------------------------------------------------------- # +# Emission record # +# --------------------------------------------------------------------------- # + + +@dataclass +class GeomEmission: + """One MuJoCo ```` to emit, in dimos world frame. + + Either ``mesh_path`` is set (mesh-type geom — also emits a + ```` asset) or the primitive parameters (``size``, ``pos``, + ``quat``) are set. + """ + + name: str + purpose: Literal["collision", "visual"] + kind: Literal["mesh", "box", "sphere", "cylinder", "capsule", "plane"] + + #: For ``kind="mesh"``: absolute path to the OBJ. + mesh_path: Path | None = None + + #: Primitive size (semantics depend on ``kind``): + #: * box → (hx, hy, hz) (half-extents) + #: * sphere → (r,) + #: * cylinder→ (r, half_height) + #: * capsule → (r, half_height) (caps extend beyond half_height) + #: * plane → (hx, hy, _grid_spacing) — last is cosmetic only + size: tuple[float, ...] | None = None + + #: World-frame position of the primitive centre. ``None`` for meshes + #: (their geometry is already in world frame). + pos: tuple[float, float, float] | None = None + + #: World-frame orientation (wxyz, MuJoCo convention). ``None`` → + #: identity. Not used for meshes. + quat: tuple[float, float, float, float] | None = None + + #: Optional friction override (slide, spin, roll). + friction: tuple[float, float, float] | None = None + + +# --------------------------------------------------------------------------- # +# Math helpers # +# --------------------------------------------------------------------------- # + + +def _matrix_to_quat_wxyz(R: np.ndarray) -> tuple[float, float, float, float]: + """3x3 right-handed rotation → quaternion ``(w, x, y, z)``. + + Standard Shepperd's method; avoids the singularity at ``trace == -1``. + """ + R = np.asarray(R, dtype=np.float64) + tr = R[0, 0] + R[1, 1] + R[2, 2] + if tr > 0: + s = 0.5 / np.sqrt(tr + 1.0) + w = 0.25 / s + x = (R[2, 1] - R[1, 2]) * s + y = (R[0, 2] - R[2, 0]) * s + z = (R[1, 0] - R[0, 1]) * s + elif (R[0, 0] > R[1, 1]) and (R[0, 0] > R[2, 2]): + s = 2.0 * np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2]) + w = (R[2, 1] - R[1, 2]) / s + x = 0.25 * s + y = (R[0, 1] + R[1, 0]) / s + z = (R[0, 2] + R[2, 0]) / s + elif R[1, 1] > R[2, 2]: + s = 2.0 * np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2]) + w = (R[0, 2] - R[2, 0]) / s + x = (R[0, 1] + R[1, 0]) / s + y = 0.25 * s + z = (R[1, 2] + R[2, 1]) / s + else: + s = 2.0 * np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1]) + w = (R[1, 0] - R[0, 1]) / s + x = (R[0, 2] + R[2, 0]) / s + y = (R[1, 2] + R[2, 1]) / s + z = 0.25 * s + return (float(w), float(x), float(y), float(z)) + + +def _quat_z_to(axis: np.ndarray) -> tuple[float, float, float, float]: + """Quaternion that rotates ``+Z`` onto ``axis`` (unit vector). + + Used for cylinder/capsule placement — MuJoCo's primitive long-axis + is local +Z; we orient by aligning Z to the prim's PCA principal + direction. + """ + z = np.array([0.0, 0.0, 1.0]) + a = axis / (np.linalg.norm(axis) + 1e-12) + d = float(np.dot(z, a)) + if d > 0.99999: + return (1.0, 0.0, 0.0, 0.0) + if d < -0.99999: + # 180° about any axis perpendicular to Z; use X. + return (0.0, 1.0, 0.0, 0.0) + cross = np.cross(z, a) + s = float(np.sqrt(2.0 * (1.0 + d))) + w = s * 0.5 + xyz = cross / s + return (w, float(xyz[0]), float(xyz[1]), float(xyz[2])) + + +# --------------------------------------------------------------------------- # +# Primitive fits # +# --------------------------------------------------------------------------- # + + +#: Lower bound on any primitive's half-extent / radius. MuJoCo rejects +#: a ```` with any size component <= 0, but truly flat prims (road +#: tiles, ceiling panels) can come out of the OBB / cylinder fit with one +#: axis at exactly 0. Clamping to 1 mm yields a valid geom that's still +#: physically reasonable as a thin slab. +_MIN_SIZE_M = 1e-3 +_SHEET_PRISM_THICKNESS_M = 0.03 +_SHEET_BOX_FILL_MIN = 0.85 +_SHEET_BOX_FILL_MAX = 1.15 +_HORIZONTAL_BOX_MAX_THICKNESS_M = 0.05 +_SHEET_PRISM_MIN_FOOTPRINT_AREA_M2 = 2.0 +_SHEET_PRISM_MAX_TRIANGLES = 1024 + + +def _fit_aabb_box(vertices: np.ndarray) -> PrimitiveFit: + """Axis-aligned bounding box. Identity quat.""" + mn, mx = vertices.min(0), vertices.max(0) + half_ext = np.maximum((mx - mn) / 2.0, _MIN_SIZE_M) + center = (mx + mn) / 2.0 + return { + "type": "box", + "size": tuple(map(float, half_ext)), + "pos": tuple(map(float, center)), + "quat": (1.0, 0.0, 0.0, 0.0), + "volume": float(np.prod(2.0 * half_ext)), + } + + +def _fit_obb_box(vertices: np.ndarray) -> PrimitiveFit: + """Oriented bounding box via PCA. Tighter than AABB when the prim + is rotated relative to world axes (most UE props are world-aligned, + so OBB ≈ AABB, but rotated assets benefit).""" + centroid = vertices.mean(0) + centered = vertices - centroid + cov = np.cov(centered.T) + _, evecs = np.linalg.eigh(cov) + # Ensure right-handed. + if np.linalg.det(evecs) < 0: + evecs[:, 0] = -evecs[:, 0] + local = centered @ evecs + mn, mx = local.min(0), local.max(0) + half_ext = np.maximum((mx - mn) / 2.0, _MIN_SIZE_M) + local_center = (mx + mn) / 2.0 + world_center = centroid + evecs @ local_center + return { + "type": "box", + "size": tuple(map(float, half_ext)), + "pos": tuple(map(float, world_center)), + "quat": _matrix_to_quat_wxyz(evecs), + "volume": float(np.prod(2.0 * half_ext)), + } + + +def _fit_sphere(vertices: np.ndarray) -> PrimitiveFit: + """Centroid + farthest-vertex. Looser than Welzl/Ritter but fine for + fill-ratio comparison.""" + centroid = vertices.mean(0) + r = max(float(np.linalg.norm(vertices - centroid, axis=1).max()), _MIN_SIZE_M) + return { + "type": "sphere", + "size": (r,), + "pos": tuple(map(float, centroid)), + "quat": (1.0, 0.0, 0.0, 0.0), + "volume": float((4.0 / 3.0) * np.pi * r**3), + } + + +def _fit_cylinder(vertices: np.ndarray) -> PrimitiveFit: + """Cylinder along PCA principal axis.""" + centroid = vertices.mean(0) + centered = vertices - centroid + cov = np.cov(centered.T) + evals, evecs = np.linalg.eigh(cov) + axis = evecs[:, -1] # largest variance + proj = centered @ axis + half_h = max(float((proj.max() - proj.min()) / 2.0), _MIN_SIZE_M) + centre_along = float((proj.max() + proj.min()) / 2.0) + pos = centroid + axis * centre_along + # radius = max perp distance from axis line + perp = centered - np.outer(centered @ axis, axis) + r = max(float(np.linalg.norm(perp, axis=1).max()), _MIN_SIZE_M) + return { + "type": "cylinder", + "size": (r, half_h), + "pos": tuple(map(float, pos)), + "quat": _quat_z_to(axis), + "volume": float(np.pi * r * r * 2.0 * half_h), + } + + +def _fit_capsule(vertices: np.ndarray) -> PrimitiveFit: + """Capsule along PCA principal axis. MuJoCo capsule half-height is + the *cylindrical* portion only; total length = 2*(half_h + r).""" + cyl = _fit_cylinder(vertices) + r, h = cyl["size"] + new_h = max(float(h - r), _MIN_SIZE_M) + vol = float(np.pi * r * r * 2.0 * new_h) + float((4.0 / 3.0) * np.pi * r**3) + return { + "type": "capsule", + "size": (r, new_h), + "pos": cyl["pos"], + "quat": cyl["quat"], + "volume": vol, + } + + +def _hull_volume(vertices: np.ndarray) -> float | None: + """Convex-hull volume in m³, or ``None`` if qhull rejects the points.""" + try: + return float(ConvexHull(vertices).volume) + except (QhullError, ValueError): + return None + + +def _mesh_volume(vertices: np.ndarray, triangles: np.ndarray) -> float: + """Signed mesh volume (Gauss / divergence theorem on triangle fans). + + Closed meshes return a positive number; for non-closed inputs the + absolute value is a coarse estimate.""" + v0 = vertices[triangles[:, 0]] + v1 = vertices[triangles[:, 1]] + v2 = vertices[triangles[:, 2]] + return float(abs(np.sum(np.einsum("ij,ij->i", v0, np.cross(v1, v2))) / 6.0)) + + +def _best_primitive_fit( + vertices: np.ndarray, + hull_vol: float, + candidates: tuple[str, ...] = ("box", "cylinder", "sphere", "capsule"), +) -> PrimitiveFit | None: + """Try every primitive in ``candidates``; return the one with the + highest fill ratio. Returns ``None`` if no fit succeeds (e.g. < 4 + points).""" + fitters = { + "box": _fit_obb_box, + "sphere": _fit_sphere, + "cylinder": _fit_cylinder, + "capsule": _fit_capsule, + } + fits: list[PrimitiveFit] = [] + for kind in candidates: + try: + f = fitters[kind](vertices) + if f["volume"] <= 0: + continue + f["fill_ratio"] = hull_vol / f["volume"] + fits.append(f) + except Exception as e: + logger.debug(f" primitive fit {kind} failed: {e}") + if not fits: + return None + return max(fits, key=lambda f: f["fill_ratio"]) + + +# --------------------------------------------------------------------------- # +# Generic geometric heuristics # +# --------------------------------------------------------------------------- # + + +def _is_tiny(extent: np.ndarray, threshold_m: float) -> bool: + return bool(extent.max() < threshold_m) + + +def _is_slab(extent: np.ndarray, aspect_ratio: float) -> bool: + """Wall / floor / door / panel — one axis is much smaller than the + other two (or one much larger than the other two — covers beams).""" + sorted_ext = np.sort(extent) + if sorted_ext[0] < 1e-6: + return True + return bool((sorted_ext[2] / sorted_ext[0]) >= aspect_ratio) + + +def _sheet_footprint_stats( + vertices: np.ndarray, + triangles: np.ndarray, + thin_axis: int, +) -> tuple[float, float] | None: + """Return ``(projected_aabb_area, projected_triangle_fill)`` for a sheet.""" + axes = [i for i in range(3) if i != thin_axis] + projected = vertices[:, axes] + span = projected.max(axis=0) - projected.min(axis=0) + box_area = float(span[0] * span[1]) + if box_area < 1e-9: + return None + + tri = projected[triangles] + edge_a = tri[:, 1] - tri[:, 0] + edge_b = tri[:, 2] - tri[:, 0] + area = 0.5 * np.abs(edge_a[:, 0] * edge_b[:, 1] - edge_a[:, 1] * edge_b[:, 0]).sum() + fill = float(area / box_area) + return box_area, fill + + +def _is_boxlike_sheet( + vertices: np.ndarray, + triangles: np.ndarray, + thin_axis: int, +) -> bool: + """Whether a thin mesh roughly fills its projected bounding rectangle. + + A single primitive box is only acceptable when the source sheet's + projected triangle area is close to the projected AABB area. Low + ratios mean an L-shape / beam strip / holes; high ratios usually mean + overlapping, folded, or angled sheets inside one prim. + """ + stats = _sheet_footprint_stats(vertices, triangles, thin_axis) + if stats is None: + return False + _, fill = stats + return _SHEET_BOX_FILL_MIN <= fill <= _SHEET_BOX_FILL_MAX + + +def _should_emit_triangle_prisms( + vertices: np.ndarray, + triangles: np.ndarray, + thin_axis: int, +) -> bool: + """Use exact-ish triangle prisms only for large horizontal sheets. + + This avoids placing huge slabs over angular floors and roof strips, + without exploding tiny decorative meshes into thousands of geoms. + """ + if thin_axis != 2: + return False + if len(triangles) > _SHEET_PRISM_MAX_TRIANGLES: + return False + stats = _sheet_footprint_stats(vertices, triangles, thin_axis) + if stats is None: + return False + footprint_area, _ = stats + return footprint_area >= _SHEET_PRISM_MIN_FOOTPRINT_AREA_M2 + + +def _thin_sheet_hulls( + vertices: np.ndarray, + triangles: np.ndarray, + thickness: float = _SHEET_PRISM_THICKNESS_M, +) -> list[tuple[np.ndarray, np.ndarray]]: + """Represent a thin non-rectangular sheet as convex triangle prisms.""" + hulls: list[tuple[np.ndarray, np.ndarray]] = [] + faces = np.asarray( + [ + [0, 1, 2], + [5, 4, 3], + [0, 3, 4], + [0, 4, 1], + [1, 4, 5], + [1, 5, 2], + [2, 5, 3], + [2, 3, 0], + ], + dtype=np.int32, + ) + + for tri_idx in triangles: + tri = vertices[tri_idx].astype(np.float64) + if not np.isfinite(tri).all(): + continue + normal = np.cross(tri[1] - tri[0], tri[2] - tri[0]) + norm = float(np.linalg.norm(normal)) + if norm < 1e-9: + continue + offset = normal / norm * (thickness * 0.5) + prism = np.vstack((tri + offset, tri - offset)).astype(np.float32) + hulls.append((prism, faces)) + + return hulls + + +def _is_flat_horizontal_box(extent: np.ndarray, thin_axis: int) -> bool: + """Thin in world Z, broad in world X/Y, and flat enough to box safely. + + PCA boxes are unstable for nearly flat floors/ceilings: any small + triangulation asymmetry can rotate the OBB basis and turn a walkable + surface into a shallow ramp. For world-horizontal slabs, the AABB is + the physically safer collision approximation. + """ + if thin_axis != 2: + return False + xy_min = float(min(extent[0], extent[1])) + z_extent = float(extent[2]) + if xy_min < 1e-6: + return False + return z_extent <= _HORIZONTAL_BOX_MAX_THICKNESS_M + + +# --------------------------------------------------------------------------- # +# Dispatcher: per-prim decision # +# --------------------------------------------------------------------------- # + + +@dataclass +class PrimDecision: + """What the dispatcher decided for one prim. Consumed by the bake + which materialises ``GeomEmission`` records and writes OBJs.""" + + #: ``"skip"`` (no collision), ``"primitive"`` (one ```` with + #: kind ∈ {box, sphere, cylinder, capsule, plane}), or ``"hulls"`` + #: (one or more mesh ````s from convex-hull decomposition). + mode: Literal["skip", "primitive", "hulls"] + + #: For ``"primitive"``: the fit dict (``type``, ``size``, ``pos``, + #: ``quat``, ``volume``, ``fill_ratio``). + primitive: PrimitiveFit | None = None + + #: For ``"hulls"``: list of ``(vertices, triangles)`` ready to write. + hulls: list[tuple[np.ndarray, np.ndarray]] = field(default_factory=list) + + #: For diagnostics: which rule fired. + reason: str = "" + + #: Optional friction override from sidecar. + friction: tuple[float, float, float] | None = None + + #: Optional per-mesh triangle cap from sidecar. The bake applies this + #: before writing OBJ assets for mesh geoms. + target_faces: int | None = None + + +def decide_for_prim( + vertices: np.ndarray, + triangles: np.ndarray, + prim_path: str, + spec: CollisionSpec, +) -> PrimDecision: + """Resolve sidecar + heuristics + auto-fit for a single prim. + + Pure function — does no I/O. The caller (bake) materialises the + decision: writes hull OBJs to disk, emits MJCF lines. + + Args: + vertices: ``(N, 3)`` float, world-frame after ``SceneMeshAlignment``. + triangles: ``(M, 3)`` int vertex indices. + prim_path: USD-style prim path used as sidecar key. + spec: parsed sidecar. + """ + extent = vertices.max(0) - vertices.min(0) + override = spec.resolve(prim_path) + kind = override.get("type", spec.default) + friction = override.get("friction") + if friction is not None: + friction = tuple(float(x) for x in friction) + target_faces = _target_faces(override) + + # 0. Explicit "skip" — short-circuit. + if kind == "skip": + return PrimDecision(mode="skip", reason="sidecar:skip", friction=friction) + + # 1. Tiny-prim guard (applies to "auto" path; explicit overrides win). + if kind in ("auto",) and _is_tiny(extent, spec.tiny_prim_extent_m): + return PrimDecision(mode="skip", reason="tiny-prim", friction=friction) + + # 2. Explicit primitive in sidecar — fit if size/pos not provided. + if kind in ("box", "sphere", "cylinder", "capsule", "plane"): + fit = _resolve_explicit_primitive(vertices, kind, override) + fit["fill_ratio"] = float("nan") # unknown — user asserted this shape + return PrimDecision( + mode="primitive", primitive=fit, reason=f"sidecar:{kind}", friction=friction + ) + + # 3. Explicit hull / decompose paths. + if kind == "hull": + return PrimDecision( + mode="hulls", + hulls=[(vertices, triangles)], # signal: single-hull, no decomp + reason="sidecar:hull", + friction=friction, + target_faces=target_faces, + ) + if kind == "decompose": + max_h = int(override.get("max_hulls", spec.coacd_max_hulls)) + hulls = _coacd_decompose(vertices, triangles, spec.coacd_threshold, max_h) + return PrimDecision( + mode="hulls", + hulls=hulls, + reason="sidecar:decompose", + friction=friction, + target_faces=target_faces, + ) + if kind in {"mesh", "decimate"}: + return PrimDecision( + mode="hulls", + hulls=[(vertices, triangles)], + reason=f"sidecar:{kind}", + friction=friction, + target_faces=target_faces, + ) + + # 4. From here on: kind == "auto". Generic heuristics first. + + # 4a. Aspect-ratio: slab/beam → force box (fill ratio may be + # marginal because of moulding/profile, but a box collision is + # the right physical answer for walls and slabs). Non-rectangular + # sheets are emitted as triangle prisms so we don't fill holes or + # angular roof/floor outlines with one huge invisible slab. + if _is_slab(extent, spec.aspect_ratio_box): + thin_axis = int(np.argmin(extent)) + if ( + spec.enable_sheet_prisms + and not _is_boxlike_sheet(vertices, triangles, thin_axis) + and _should_emit_triangle_prisms(vertices, triangles, thin_axis) + ): + hulls = _thin_sheet_hulls(vertices, triangles) + if hulls: + return PrimDecision( + mode="hulls", + hulls=hulls, + reason=f"thin-sheet:triangle-prisms({len(hulls)})", + friction=friction, + ) + + if _is_flat_horizontal_box(extent, thin_axis): + fit = _fit_aabb_box(vertices) + reason = "aspect-ratio:horizontal-slab" + else: + fit = _fit_obb_box(vertices) + reason = "aspect-ratio:slab" + fit["fill_ratio"] = float("nan") + return PrimDecision(mode="primitive", primitive=fit, reason=reason, friction=friction) + + # 4b. Need hull volume for the rest. + hull_vol = _hull_volume(vertices) + if hull_vol is None: + return PrimDecision(mode="skip", reason="degenerate (qhull rejected)", friction=friction) + + # 4c. Try primitive auto-fit. + auto_fit = _best_primitive_fit(vertices, hull_vol) + if auto_fit is not None and 0.0 < auto_fit["fill_ratio"] <= 1.5: + # fill_ratio > 1 happens for non-closed hulls; cap to keep this + # finite when reporting. Accept if within tolerance. + if auto_fit["fill_ratio"] >= spec.fill_threshold: + return PrimDecision( + mode="primitive", + primitive=auto_fit, + reason=f"auto:{auto_fit['type']}({auto_fit['fill_ratio']:.2f})", + friction=friction, + ) + + # 4d. Near-convex shortcut: skip CoACD, single hull. + mesh_vol = _mesh_volume(vertices, triangles) + if hull_vol > 0 and mesh_vol / hull_vol > spec.near_convex_threshold: + return PrimDecision( + mode="hulls", + hulls=[(vertices, triangles)], + reason=f"near-convex({mesh_vol / hull_vol:.2f})", + friction=friction, + ) + + # 4e. Small concave prim → single hull (matches today's behaviour + # for furniture-scale things; faster than CoACD). + if hull_vol < spec.shell_volume_m3: + return PrimDecision( + mode="hulls", + hulls=[(vertices, triangles)], + reason="small-shell:single-hull", + friction=friction, + ) + + # 4f. Large concave shell → CoACD. + hulls = _coacd_decompose(vertices, triangles, spec.coacd_threshold, spec.coacd_max_hulls) + return PrimDecision( + mode="hulls", + hulls=hulls, + reason=f"coacd:{len(hulls)}", + friction=friction, + ) + + +# --------------------------------------------------------------------------- # +# Helpers used by the dispatcher # +# --------------------------------------------------------------------------- # + + +def _resolve_explicit_primitive( + vertices: np.ndarray, + kind: str, + override: OverrideConfig, +) -> PrimitiveFit: + """Build a primitive fit dict from a sidecar override. + + If the override supplies ``size`` (and optionally ``pos`` / ``quat``), + those win. Otherwise we auto-fit the requested primitive and use + those params. ``plane`` is special-cased — we always derive from + the prim's xy footprint at its min z. + """ + if kind == "plane": + mn = vertices.min(0) + mx = vertices.max(0) + return { + "type": "plane", + "size": (float((mx[0] - mn[0]) / 2.0), float((mx[1] - mn[1]) / 2.0), 0.5), + "pos": ( + float((mx[0] + mn[0]) / 2.0), + float((mx[1] + mn[1]) / 2.0), + float(mn[2]), + ), + "quat": (1.0, 0.0, 0.0, 0.0), + "volume": 0.0, + } + + box_fitter = _fit_aabb_box if "min_thickness" in override else _fit_obb_box + fitters = { + "box": box_fitter, + "sphere": _fit_sphere, + "cylinder": _fit_cylinder, + "capsule": _fit_capsule, + } + fit = fitters[kind](vertices) + # Apply explicit overrides if provided. + if "size" in override: + fit["size"] = tuple(float(x) for x in override["size"]) + if "pos" in override: + fit["pos"] = tuple(float(x) for x in override["pos"]) + if "quat" in override: + fit["quat"] = tuple(float(x) for x in override["quat"]) + if kind == "box": + _apply_box_min_thickness(fit, vertices, override) + return fit + + +def _apply_box_min_thickness( + fit: PrimitiveFit, + vertices: np.ndarray, + override: OverrideConfig, +) -> None: + raw_min_thickness = override.get("min_thickness") + if raw_min_thickness is None: + return + + min_half_z = max(float(raw_min_thickness) * 0.5, _MIN_SIZE_M) + size = np.asarray(fit["size"], dtype=np.float64) + pos = np.asarray(fit["pos"], dtype=np.float64) + if size.shape[0] < 3 or pos.shape[0] < 3 or size[2] >= min_half_z: + return + + old_half_z = float(size[2]) + size[2] = min_half_z + + preserve = str(override.get("preserve", "center")).lower() + if preserve == "top": + top_z = ( + pos[2] + old_half_z + if "pos" in override or "size" in override + else float(np.max(vertices[:, 2])) + ) + pos[2] = top_z - min_half_z + elif preserve == "bottom": + bottom_z = ( + pos[2] - old_half_z + if "pos" in override or "size" in override + else float(np.min(vertices[:, 2])) + ) + pos[2] = bottom_z + min_half_z + elif preserve in {"center", "centre"}: + pass + else: + raise ValueError("box min_thickness preserve must be one of: top, bottom, center") + + fit["size"] = tuple(map(float, size)) + fit["pos"] = tuple(map(float, pos)) + + +def _target_faces(override: OverrideConfig) -> int | None: + raw = override.get("target_faces", override.get("max_faces")) + if raw is None: + return None + target_faces = int(raw) + if target_faces <= 0: + return None + return max(4, target_faces) + + +def _coacd_decompose( + vertices: np.ndarray, + triangles: np.ndarray, + threshold: float, + max_hulls: int, +) -> list[tuple[np.ndarray, np.ndarray]]: + """Run CoACD on a single prim, return list of ``(verts, tris)`` hulls. + + CoACD is imported lazily — it ships its own C library and we don't + want every import of ``collision_spec`` to pay that cost. + """ + import coacd # type: ignore[import-not-found, import-untyped] + + # CoACD's C lib prints a lot per invocation; quiet it once per process. + if not getattr(_coacd_decompose, "_silenced", False): + coacd.set_log_level("error") + _coacd_decompose._silenced = True # type: ignore[attr-defined] + + mesh = coacd.Mesh(vertices.astype(np.float64), triangles.astype(np.int32)) + # CoACD's MCTS defaults (mcts_iterations=150, resolution=2000) are tuned + # for offline asset prep — minutes per shell on a multi-thousand-prim + # scene. We dial both down ~5x; the resulting hulls are slightly + # noisier but the bake finishes in minutes, not hours. For a one-off + # final bake users can override via the sidecar (future work) or call + # ``bake_scene_mjcf`` directly with a custom ``CollisionSpec``. + parts = coacd.run_coacd( + mesh, + threshold=threshold, + max_convex_hull=max_hulls, + resolution=500, + mcts_iterations=30, + mcts_nodes=10, + ) + out: list[tuple[np.ndarray, np.ndarray]] = [] + for v, t in parts: + v = np.asarray(v, dtype=np.float32) + t = np.asarray(t, dtype=np.int32) + if len(v) >= 4 and len(t) >= 1: + out.append((v, t)) + return out + + +__all__ = [ + "CollisionSpec", + "GeomEmission", + "PrimDecision", + "decide_for_prim", +] diff --git a/dimos/simulation/mujoco/scene_mesh_to_mjcf.py b/dimos/simulation/mujoco/scene_mesh_to_mjcf.py new file mode 100644 index 0000000000..67af2a04c2 --- /dev/null +++ b/dimos/simulation/mujoco/scene_mesh_to_mjcf.py @@ -0,0 +1,1003 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Bake a scene mesh into an MJCF wrapper around a robot MJCF. + +The bake walks every prim returned by :func:`load_scene_prims` and asks +:func:`dimos.simulation.mujoco.collision_spec.decide_for_prim` what to +emit for it. The dispatcher returns one of three modes: + +- ``"primitive"`` -- a single MuJoCo primitive ```` (box / sphere / + cylinder / capsule / plane). Used when the prim is approximately + prismatic (auto-fit) or when a sidecar override forces it. +- ``"hulls"`` -- one or more mesh ````s. Either a single convex + hull (small / near-convex prims) or a CoACD decomposition (genuine + concave shells: stairs, planters). +- ``"skip"`` -- no collision geom at all. Used for sidecar-tagged + decoration (lamps, signs) and prims smaller than a threshold. + +Hulls produced by either path are validated with :func:`_valid_hull` +(matrix-rank coplanarity check + scipy ``Qt`` qhull pre-flight); when a +hull is invalid we fall back to a thin OBB box via +:func:`_fallback_box_geom` rather than dropping the geometry, so the +robot doesn't sink through holes. + +When ``include_visual_mesh=True`` the bake additionally writes the +prim's original triangles as a non-colliding visual geom (group 2, +``contype=0 conaffinity=0``). UE's USD exporter culls hidden faces on +static meshes (a floor slab ships with only top + bottom face pairs, +no sides) -- we route visual writes through :func:`_write_visual_obj`, +which substitutes the prim's convex hull when it isn't watertight, so +the viewer renders solid geometry instead of see-through slabs. + +Per-prim work is fanned across worker processes since each prim's +decision is independent and CoACD calls dominate wall time. Standalone +CLI bakes use forked processes; in an already-threaded DimOS runtime we +use ``forkserver`` so workers do not inherit the parent process's active +threads. + +Output is cached at ``~/.cache/dimos/scene_meshes//`` keyed on +the SHA256 of (source mesh, robot MJCF, alignment, meshdir, sidecar +spec, visual flag, schema version). :func:`load_or_bake` is the +recommended entry point -- it handles a three-tier cache: + + 1. ``compiled.mjb`` exists -> load directly (~1 s) + 2. ``wrapper.xml`` + OBJs exist -> compile XML, save ``.mjb`` + 3. Nothing exists -> full bake, then compile + save ``.mjb`` +""" + +from __future__ import annotations + +import argparse +from concurrent.futures import ProcessPoolExecutor, as_completed +from dataclasses import asdict, dataclass, replace +import hashlib +import multiprocessing +import os +from pathlib import Path +import time +from typing import Any + +import numpy as np +import open3d as o3d # type: ignore[import-untyped] + +from dimos.simulation.mujoco.collision_spec import ( + CollisionSpec, + decide_for_prim, +) +from dimos.simulation.scene_assets.mesh_scene import ( + SceneMeshAlignment, + ScenePrimMesh, + load_scene_prims, + split_disconnected_scene_prims, +) +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + + +CACHE_DIR = Path.home() / ".cache" / "dimos" / "scene_meshes" + + +# Scene-only wrapper -- no robot include. Robots are attached at runtime +# via ``MjSpec.attach()`` (see ``MujocoSimModule.start``), keeping the +# cooked package robot-agnostic. ``meshdir="."`` resolves the cooked +# scene OBJs that sit alongside this file. +# +# The dummy ```` on dimos_scene bypasses MuJoCo's auto- +# computation of body inertia from geom volumes -- the body is static +# (no joint) so the values don't affect physics, but without this any +# zero-volume visual mesh (road tiles, ceiling panels, flat slabs) +# triggers ``Error: mesh volume is too small`` at compile time. +_WRAPPER_TEMPLATE = """\ + + + + + + + +{asset_meshes} + + + + +{scene_geoms} + + + +""" + +# ``inertia="shell"`` makes MuJoCo compute mesh inertia from surface +# area instead of enclosed volume -- robust to non-watertight visual +# meshes from art tools. Safe for closed CoACD hulls too, so we apply +# it universally for one fewer code path. +_ASSET_LINE = ' ' + +# Collision (group 3) -- actually collides. Keep it opaque so MuJoCo +# depth renders treat the scene as solid for lidar/camera simulation. +_COL_MESH_LINE = ( + ' ' +) +_COL_BOX_LINE = ( + ' ' +) +_COL_SPHERE_LINE = ( + ' ' +) +_COL_CYL_LINE = ( + ' ' +) +_COL_CAP_LINE = ( + ' ' +) +_COL_PLANE_LINE = ( + ' ' +) + +# Visual (group 2) -- drawn, doesn't collide. +_VISUAL_GEOM_LINE = ( + ' ' +) + + +# Constants kept from the prior implementation -- conservative +# fallback thresholds for hull validity / box-fallback geometry. +_DEGENERATE_EPS = 1e-3 +_MIN_HULL_EXTENT_M = 5e-3 +_FALLBACK_BOX_THICKNESS_M = 0.03 +_MIN_FALLBACK_BOX_EXTENT_M = 0.25 +_MIN_FALLBACK_BOX_AREA_M2 = 0.05 + +_CACHE_KEY_LEN = 12 +# Bump when the bake pipeline's output format changes so old caches +# invalidate on the next call. Increment for any change that could +# affect MJCF emission (new geom kinds, rewritten visual policy, etc.). +# This is only a local cache salt; it is not a persisted file format +# contract and old cache directories can safely stay on disk. +_CACHE_SCHEMA_VERSION = "scene-only-v10" + + +@dataclass +class _BakeArtifacts: + """Aggregated stats + emission lines from one bake.""" + + asset_lines: list[str] + geom_lines: list[str] + n_primitive: int + n_hulls_total: int + n_box_fallbacks: int + n_skipped: int + n_visuals: int + n_degenerate_dropped: int + decision_reasons: dict[str, int] + + +def _resolve_existing_file(path: str | Path, label: str) -> Path: + resolved = Path(path).expanduser().resolve() + if not resolved.exists(): + raise FileNotFoundError(f"{label} not found: {resolved}") + if not resolved.is_file(): + raise ValueError(f"{label} must be a file, got: {resolved}") + return resolved + + +def bake_scene_mjcf( + scene_mesh_path: str | Path, + alignment: SceneMeshAlignment | None = None, + cache_root: Path | None = None, + collision_spec: CollisionSpec | None = None, + include_visual_mesh: bool = False, + rebake: bool = False, +) -> Path: + """Convert ``scene_mesh_path`` to OBJs + scene-only MJCF wrapper. + + The wrapper is robot-agnostic: it declares the cooked scene as the + ``dimos_scene`` static body and nothing else. Robots are attached at + runtime via ``MjSpec.attach()`` inside ``MujocoSimModule.start``. + + Args: + scene_mesh_path: ``.usdz`` / ``.usda`` / ``.glb`` / ``.obj`` / + etc. Anything ``mesh_scene.load_scene_prims`` accepts. + alignment: scale / translation / rotation / y-up swap to bake + into world frame before any geom is emitted. + cache_root: override the cache root (defaults to + ``~/.cache/dimos/scene_meshes``). + collision_spec: per-prim policy. ``None`` auto-discovers a + sidecar ``.collision.json`` next to the source, or + falls back to ``CollisionSpec()`` defaults (auto-fit + primitives, CoACD on large concave shells). + include_visual_mesh: also emit a non-colliding visual geom for + every prim showing its original triangles. The viewer + renders these instead of the collision hulls -- much nicer + for visual debugging, but doubles disk usage. When ``True`` + non-watertight prim meshes are substituted with their convex + hull so they don't appear see-through. + rebake: ignore an existing ``wrapper.xml`` in the computed cache + directory and regenerate the scene collision geometry. + + Returns: + Path to the scene-only wrapper MJCF. Load with + ``mujoco.MjSpec.from_file`` and attach a robot via ``attach()``. + """ + scene_mesh_path = _resolve_existing_file(scene_mesh_path, "scene mesh") + align = alignment or SceneMeshAlignment() + spec = collision_spec or CollisionSpec.auto_discover(scene_mesh_path) + + cache_key = _cache_key( + scene_mesh_path, + align, + spec=spec, + include_visual_mesh=include_visual_mesh, + ) + root = (cache_root or CACHE_DIR).expanduser() + cache_dir = root / cache_key + wrapper_path = cache_dir / "wrapper.xml" + + if not rebake and _cache_hit(wrapper_path): + logger.info(f"bake_scene_mjcf: cache hit at {cache_dir}") + return wrapper_path + + cache_dir.mkdir(parents=True, exist_ok=True) + + logger.info(f"bake_scene_mjcf: loading + aligning {scene_mesh_path}") + prims = load_scene_prims(scene_mesh_path, alignment=align) + if spec.split_disconnected_components: + prims, split_stats = split_disconnected_scene_prims( + prims, + min_components=spec.split_min_components, + extent_ratio=spec.split_extent_ratio, + prim_min_extent=spec.split_prim_min_extent_m, + axis_ratio=spec.split_axis_ratio, + min_component_extent=spec.split_component_min_extent_m, + min_component_faces=spec.split_component_min_faces, + can_split=lambda prim: ( + spec.resolve(prim.prim_path or prim.name).get("type", spec.default) == "auto" + ), + ) + if split_stats["split_prims"]: + logger.info( + "bake_scene_mjcf: split %s disconnected prims into %s kept " + "components; dropped %s tiny components", + split_stats["split_prims"], + split_stats["emitted_components"], + split_stats["dropped_components"], + ) + logger.info(f"bake_scene_mjcf: {len(prims)} prims to process") + if spec.enable_sheet_prisms and len(prims) > spec.sheet_prism_max_scene_prims: + logger.info( + "bake_scene_mjcf: disabling thin-sheet triangle prisms for " + f"{len(prims)}-prim scene; use a collision sidecar to opt in" + ) + spec = replace(spec, enable_sheet_prisms=False) + scene_center, scene_extent = _scene_bounds(prims) + + artifacts = _bake_prims( + prims, + cache_dir, + spec=spec, + include_visual_mesh=include_visual_mesh, + ) + if not artifacts.geom_lines: + raise RuntimeError( + "bake_scene_mjcf: every prim got skipped or produced only " + "degenerate hulls; nothing left to collide against. Check " + "the source mesh and alignment." + ) + + logger.info( + f"bake_scene_mjcf: {artifacts.n_primitive} primitive geoms, " + f"{artifacts.n_hulls_total} hull geoms, " + f"{artifacts.n_box_fallbacks} box fallbacks, " + f"{artifacts.n_visuals} visual passthrough meshes, " + f"{artifacts.n_skipped} skipped, " + f"{artifacts.n_degenerate_dropped} degenerate hulls dropped" + ) + # Top-10 decision reasons -- useful when tuning a sidecar. + for reason, n in sorted(artifacts.decision_reasons.items(), key=lambda kv: -kv[1])[:10]: + logger.info(f" reason {reason:32s} {n}") + + _write_wrapper( + wrapper_path=wrapper_path, + cache_key=cache_key, + asset_lines=artifacts.asset_lines, + geom_lines=artifacts.geom_lines, + statistic_center=scene_center, + statistic_extent=scene_extent, + ) + return wrapper_path + + +def load_or_bake( + scene_mesh_path: str | Path, + alignment: SceneMeshAlignment | None = None, + cache_root: Path | None = None, + collision_spec: CollisionSpec | None = None, + include_visual_mesh: bool = False, + rebake: bool = False, +) -> Path: + """Return the cached or freshly baked scene-only wrapper MJCF. + + Robots are attached at runtime via ``MjSpec``; no ``compiled.mjb`` is + produced at cook time. The cache key is over the source mesh, + alignment, collision spec, and schema version -- robot-agnostic. + """ + scene_mesh_path = _resolve_existing_file(scene_mesh_path, "scene mesh") + return bake_scene_mjcf( + scene_mesh_path=scene_mesh_path, + alignment=alignment, + cache_root=cache_root, + collision_spec=collision_spec, + include_visual_mesh=include_visual_mesh, + rebake=rebake, + ) + + +# --------------------------------------------------------------------------- # +# Cache key # +# --------------------------------------------------------------------------- # + + +def _cache_key( + scene_mesh_path: Path, + alignment: SceneMeshAlignment, + *, + spec: CollisionSpec, + include_visual_mesh: bool, +) -> str: + """SHA256-12 over every input that affects bake output. + + Robot-agnostic: the cooked scene wrapper is the same regardless of + which robot will eventually be attached at runtime via ``MjSpec``. + """ + import json + + def _file_signature(path: Path) -> str: + st = path.stat() + return f"{path}:{st.st_size}:{st.st_mtime_ns}" + + h = hashlib.sha256() + h.update(_CACHE_SCHEMA_VERSION.encode()) + h.update(_file_signature(scene_mesh_path).encode()) + h.update(repr(sorted(asdict(alignment).items())).encode()) + h.update(json.dumps(asdict(spec), sort_keys=True).encode()) + h.update(b"visual=" + (b"1" if include_visual_mesh else b"0")) + return h.hexdigest()[:_CACHE_KEY_LEN] + + +def _cache_hit(wrapper_path: Path) -> bool: + if not wrapper_path.exists(): + return False + try: + text = wrapper_path.read_text() + except OSError: + return False + return " _BakeArtifacts: + """Fan per-prim work across cores; aggregate the resulting MJCF lines. + + Standalone bakes use ``fork`` so workers inherit the parent's + already-imported modules. Runtime bakes inside DimOS may happen + after other modules have started threads; in that case use + ``forkserver`` so workers do not inherit locks from the parent + process's C extension state. + """ + asset_lines: list[str] = [] + geom_lines: list[str] = [] + n_primitive = 0 + n_hulls_total = 0 + n_box_fallbacks = 0 + n_skipped = 0 + n_visuals = 0 + n_degenerate = 0 + reasons: dict[str, int] = {} + + work_items = [(prim, cache_dir, spec, include_visual_mesh) for prim in prims] + n_workers = max(1, (os.cpu_count() or 4) - 1) + if _native_thread_count() > 1: + n_workers = min(n_workers, 8) + start_method = ( + "forkserver" if "forkserver" in multiprocessing.get_all_start_methods() else "spawn" + ) + else: + start_method = "fork" + logger.info( + f"_bake_prims: fanning {len(prims)} prims across {n_workers} workers ({start_method})" + ) + + t0 = time.time() + mp_ctx = multiprocessing.get_context(start_method) + executor = ProcessPoolExecutor(max_workers=n_workers, mp_context=mp_ctx) + + progress_every = 25 if len(prims) <= 500 else 250 + with executor as ex: + futures = [ex.submit(_process_one_prim, item) for item in work_items] + done = 0 + for fut in as_completed(futures): + a_lines, g_lines, mode, reason, counters = fut.result() + asset_lines.extend(a_lines) + geom_lines.extend(g_lines) + reasons[reason] = reasons.get(reason, 0) + 1 + if mode == "primitive": + n_primitive += 1 + elif mode == "skip": + n_skipped += 1 + n_hulls_total += counters["hulls"] + n_box_fallbacks += counters["box_fallbacks"] + n_visuals += counters["visuals"] + n_degenerate += counters["degenerate"] + done += 1 + if done % progress_every == 0 or done == len(prims): + elapsed = time.time() - t0 + eta = elapsed * (len(prims) - done) / max(done, 1) + logger.info( + f" prim {done}/{len(prims)} " + f"({100 * done / len(prims):.0f}%) " + f"elapsed={elapsed:.0f}s eta={eta:.0f}s " + f"hulls_so_far={n_hulls_total}" + ) + + return _BakeArtifacts( + asset_lines=asset_lines, + geom_lines=geom_lines, + n_primitive=n_primitive, + n_hulls_total=n_hulls_total, + n_box_fallbacks=n_box_fallbacks, + n_skipped=n_skipped, + n_visuals=n_visuals, + n_degenerate_dropped=n_degenerate, + decision_reasons=reasons, + ) + + +def _native_thread_count() -> int: + try: + return len(os.listdir("/proc/self/task")) + except OSError: + return 1 + + +# --------------------------------------------------------------------------- # +# Geom emission helpers # +# --------------------------------------------------------------------------- # + + +def _emit_primitive_geom( + prim_name: str, + fit: dict[str, Any] | None, + friction_attr: str, +) -> str | None: + """Render one ``PrimDecision.primitive`` dict to MJCF text. + + Returns ``None`` if ``fit`` is missing required fields (defensive -- + ``decide_for_prim`` should always populate them, but a malformed + sidecar override could slip through). + """ + if fit is None: + return None + kind = fit.get("type") + pos = _fmt_vec(np.asarray(fit["pos"])) + size = _fmt_vec(np.asarray(fit["size"])) + quat = ( + _fmt_vec(np.asarray(fit["quat"])) + if "quat" in fit and fit["quat"] is not None + else "1 0 0 0" + ) + name = f"{prim_name}_col" + if kind == "box": + return _COL_BOX_LINE.format( + name=name, pos=pos, quat=quat, size=size, friction=friction_attr + ) + if kind == "sphere": + return _COL_SPHERE_LINE.format(name=name, pos=pos, size=size, friction=friction_attr) + if kind == "cylinder": + return _COL_CYL_LINE.format( + name=name, pos=pos, quat=quat, size=size, friction=friction_attr + ) + if kind == "capsule": + return _COL_CAP_LINE.format( + name=name, pos=pos, quat=quat, size=size, friction=friction_attr + ) + if kind == "plane": + return _COL_PLANE_LINE.format( + name=name, pos=pos, quat=quat, size=size, friction=friction_attr + ) + return None + + +# --------------------------------------------------------------------------- # +# Hull validity & box fallback (preserved from prior implementation) # +# --------------------------------------------------------------------------- # + + +def _valid_hull(v: np.ndarray, f: np.ndarray) -> bool: + """Reject hulls that MuJoCo's qhull would choke on at compile time. + + Four layers: + 1. trivial -- < 4 vertices or < 4 faces. + 2. extent -- all-axis ``> 5 mm`` (matches MuJoCo's mj_loadXML + coplanarity tolerance for ~100mm-wide hulls). + 3. rank -- centred vertex matrix must have rank 3 (catches + coplanar hulls the extent check misses, e.g. a T-shaped + hull whose XY extent is large but Z is zero). + 4. scipy ConvexHull pre-flight with ``Qt`` -- same options + MuJoCo uses; if scipy can't build it, mj_loadXML can't either. + """ + if len(v) < 4 or len(f) < 4: + return False + extent = v.max(axis=0) - v.min(axis=0) + if (extent < _DEGENERATE_EPS).any(): + return False + if float(extent.min()) < _MIN_HULL_EXTENT_M: + return False + centered = v.astype(np.float64) - v.astype(np.float64).mean(axis=0) + if np.linalg.matrix_rank(centered, tol=_DEGENERATE_EPS) < 3: + return False + try: + from scipy.spatial import ConvexHull, QhullError # type: ignore[import-untyped] + + ConvexHull(v, qhull_options="Qt") + except (QhullError, ValueError): + return False + return True + + +def _fallback_box_geom(name: str, vertices: np.ndarray, friction_attr: str = "") -> str | None: + """Emit a thin OBB box geom for vertices that can't form a valid hull. + + The thickness floor (``_FALLBACK_BOX_THICKNESS_M = 3 cm``) keeps the + box thick enough that the robot can stand on it without falling + through. Returns ``None`` for prims too small to bother (< 25 cm + largest extent or < 0.05 m^2 face area) -- those fall through to + the degenerate counter. + """ + finite = vertices[np.isfinite(vertices).all(axis=1)].astype(np.float64) + if len(finite) < 3: + return None + aabb_extent = finite.max(axis=0) - finite.min(axis=0) + sorted_extents = np.sort(aabb_extent) + if sorted_extents[-1] < _MIN_FALLBACK_BOX_EXTENT_M: + return None + if sorted_extents[-1] * sorted_extents[-2] < _MIN_FALLBACK_BOX_AREA_M2: + return None + + center, rotation, extent = _oriented_box(finite) + extent = np.maximum(extent, _FALLBACK_BOX_THICKNESS_M) + half_size = 0.5 * extent + quat = _rotation_matrix_to_wxyz(rotation) + return _COL_BOX_LINE.format( + name=name, + pos=_fmt_vec(center), + quat=_fmt_vec(quat), + size=_fmt_vec(half_size), + friction=friction_attr, + ) + + +def _oriented_box( + vertices: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """OBB via trimesh's ``bounding_box_oriented``. + + Falls back to AABB if trimesh's OBB fitter produces non-finite + output or the prim has < 3 vertices. + """ + try: + import trimesh # type: ignore[import-untyped] + + tm = trimesh.Trimesh(vertices=vertices, faces=np.empty((0, 3), dtype=np.int32)) + obb = tm.bounding_box_oriented + transform = np.asarray(obb.primitive.transform, dtype=np.float64) + extent = np.asarray(obb.primitive.extents, dtype=np.float64) + rotation = transform[:3, :3] + center = transform[:3, 3] + if np.linalg.det(rotation) < 0: + rotation[:, 0] *= -1.0 + if np.isfinite(center).all() and np.isfinite(rotation).all() and np.isfinite(extent).all(): + return center, rotation, np.abs(extent) + except Exception: + pass + + lo = vertices.min(axis=0) + hi = vertices.max(axis=0) + return (lo + hi) * 0.5, np.eye(3), hi - lo + + +def _rotation_matrix_to_wxyz(rotation: np.ndarray) -> np.ndarray: + """3x3 rotation -> ``(w, x, y, z)`` quaternion.""" + from scipy.spatial.transform import Rotation # type: ignore[import-untyped] + + xyzw = Rotation.from_matrix(rotation).as_quat() + return np.array([xyzw[3], xyzw[0], xyzw[1], xyzw[2]], dtype=np.float64) + + +def _fmt_vec(values: np.ndarray) -> str: + return " ".join(f"{float(v):.9g}" for v in values) + + +def _scene_bounds(prims: list[ScenePrimMesh]) -> tuple[np.ndarray, float]: + """Return a viewer-friendly center and extent for the aligned scene. + + MuJoCo's viewer uses ``statistic.center`` / ``statistic.extent`` for + camera framing and clipping. The included robot MJCF's defaults are + much too small for baked building-scale scenes, so wrappers need to + advertise the scene bounds explicitly. + """ + mins: list[np.ndarray] = [] + maxs: list[np.ndarray] = [] + for prim in prims: + vertices = np.asarray(prim.vertices, dtype=np.float64) + if vertices.ndim != 2 or vertices.shape[1] != 3 or len(vertices) == 0: + continue + finite = vertices[np.isfinite(vertices).all(axis=1)] + if len(finite) == 0: + continue + mins.append(finite.min(axis=0)) + maxs.append(finite.max(axis=0)) + + if not mins: + return np.zeros(3, dtype=np.float64), 1.0 + + scene_min = np.min(np.vstack(mins), axis=0) + scene_max = np.max(np.vstack(maxs), axis=0) + center = (scene_min + scene_max) * 0.5 + diagonal = scene_max - scene_min + extent = max(float(np.linalg.norm(diagonal) * 0.5 * 1.1), 1.0) + return center, extent + + +# --------------------------------------------------------------------------- # +# OBJ I/O # +# --------------------------------------------------------------------------- # + + +def _write_hull_obj(obj_file: Path, vertices: np.ndarray, faces: np.ndarray) -> None: + """Write a CoACD/single-hull mesh. No watertight check -- hulls are + closed by construction.""" + _write_mesh_obj(obj_file, vertices, faces) + + +def _simplify_mesh_geom( + vertices: np.ndarray, + faces: np.ndarray, + target_faces: int, +) -> tuple[np.ndarray, np.ndarray]: + if target_faces <= 0 or len(faces) <= target_faces: + return vertices, faces + + mesh = o3d.geometry.TriangleMesh() + mesh.vertices = o3d.utility.Vector3dVector(vertices.astype(np.float64)) + mesh.triangles = o3d.utility.Vector3iVector(faces.astype(np.int32)) + try: + mesh.remove_duplicated_vertices() + mesh.remove_duplicated_triangles() + mesh.remove_degenerate_triangles() + mesh.remove_unreferenced_vertices() + simplified = mesh + for _ in range(3): + if len(simplified.triangles) <= target_faces: + break + simplified = simplified.simplify_quadric_decimation( + target_number_of_triangles=target_faces + ) + simplified.remove_degenerate_triangles() + simplified.remove_duplicated_triangles() + simplified.remove_unreferenced_vertices() + out_vertices = np.asarray(simplified.vertices, dtype=np.float32) + out_faces = np.asarray(simplified.triangles, dtype=np.int32) + if len(out_vertices) >= 4 and 4 <= len(out_faces) <= target_faces: + return out_vertices, out_faces + except Exception: + logger.debug("mesh simplification failed; falling back to convex hull", exc_info=True) + + hull = _convex_hull_mesh(vertices) + return hull if hull is not None else (vertices, faces) + + +def _convex_hull_mesh(vertices: np.ndarray) -> tuple[np.ndarray, np.ndarray] | None: + try: + from scipy.spatial import ConvexHull, QhullError # type: ignore[import-untyped] + + hull = ConvexHull(vertices.astype(np.float64)) + except (QhullError, ValueError): + return None + + faces = np.asarray(hull.simplices, dtype=np.int32) + used = np.unique(faces.reshape(-1)) + remap = {int(old): idx for idx, old in enumerate(used)} + remapped_faces = np.vectorize(remap.__getitem__, otypes=[np.int32])(faces) + return vertices[used].astype(np.float32), remapped_faces.astype(np.int32) + + +def _write_visual_obj(obj_file: Path, vertices: np.ndarray, faces: np.ndarray) -> None: + """Write a *renderable* OBJ -- closed under all viewing angles. + + UE's static-mesh exporter culls hidden faces (a floor slab ships + with only top + bottom face pairs, no sides), so writing the + artist's geometry verbatim produces meshes that appear see-through + in MuJoCo's viewer from any oblique angle. We check + ``trimesh.is_watertight`` and, if not, substitute the prim's + convex hull (which is always closed). + + For non-prismatic prims (chairs, plants) the hull is a coarse + visual approximation; for the most common offenders (floor / roof + / wall / ceiling slabs that are box-shaped to begin with) the hull + matches the original exactly. Watertight prims (full furniture + meshes from UE) keep their original geometry. + """ + import trimesh # type: ignore[import-untyped] + + tm = trimesh.Trimesh( + vertices=np.asarray(vertices, dtype=np.float64), + faces=np.asarray(faces, dtype=np.int32), + process=False, + ) + if not tm.is_watertight: + try: + hull = tm.convex_hull + if len(hull.vertices) >= 4 and len(hull.faces) >= 4: + vertices = np.asarray(hull.vertices, dtype=np.float64) + faces = np.asarray(hull.faces, dtype=np.int32) + except Exception: + pass # fall back to original; visual may look hollow + _write_mesh_obj(obj_file, vertices, faces) + + +def _write_mesh_obj(obj_file: Path, vertices: np.ndarray, faces: np.ndarray) -> None: + o3d_mesh = o3d.geometry.TriangleMesh() + o3d_mesh.vertices = o3d.utility.Vector3dVector(vertices.astype(np.float64)) + o3d_mesh.triangles = o3d.utility.Vector3iVector(faces) + o3d_mesh.compute_vertex_normals() + if not o3d.io.write_triangle_mesh( + str(obj_file), + o3d_mesh, + write_vertex_normals=True, + write_vertex_colors=False, + ): + raise RuntimeError(f"open3d failed to write OBJ: {obj_file}") + + +# --------------------------------------------------------------------------- # +# Wrapper writer + CLI # +# --------------------------------------------------------------------------- # + + +def _write_wrapper( + *, + wrapper_path: Path, + cache_key: str, + asset_lines: list[str], + geom_lines: list[str], + statistic_center: np.ndarray, + statistic_extent: float, +) -> None: + """Emit the scene-only wrapper.xml. Robots attach at runtime via + ``MjSpec``; the wrapper directory holds only this file plus the + cooked scene OBJs that it references with relative paths. + """ + visual_zfar = max(float(statistic_extent) * 20.0, 10000.0) + wrapper_xml = _WRAPPER_TEMPLATE.format( + model_name=f"scene_{cache_key}", + statistic_center=_fmt_vec(statistic_center), + statistic_extent=f"{float(statistic_extent):.9g}", + visual_zfar=f"{visual_zfar:.9g}", + asset_meshes="\n".join(asset_lines), + scene_geoms="\n".join(geom_lines), + ) + wrapper_path.write_text(wrapper_xml) + logger.info(f"_write_wrapper: wrote {wrapper_path}") + + +def cli_main() -> None: + """``python -m dimos.simulation.mujoco.scene_mesh_to_mjcf [opts]``. + + Bake (or load from cache), optionally launch the MuJoCo viewer. + """ + p = argparse.ArgumentParser( + prog="python -m dimos.simulation.mujoco.scene_mesh_to_mjcf", + description="Bake a USD/GLB/OBJ scene into a robot-agnostic scene-only MJCF wrapper.", + ) + p.add_argument("scene", type=Path, help="scene mesh path (.usda, .usdz, .glb, ...)") + p.add_argument( + "--scale", + type=float, + default=1.0, + help="multiplicative scale (use 0.01 for UE / centimeter sources). Default 1.0.", + ) + p.add_argument( + "--no-y-up", + action="store_true", + help="source is already Z-up (UE exports with metersPerUnit=0.01 and " + "upAxis=Z). Default assumes Y-up source (Blender, glTF, Apple USDZ).", + ) + p.add_argument( + "--collision-spec", + type=Path, + default=None, + help="path to a collision-spec sidecar JSON. Default auto-discovers " + "``.collision.json`` next to the source.", + ) + p.add_argument( + "--visual", + action="store_true", + help="emit visual passthrough meshes (group 2). Off by default -- " + "saves disk and render cost, but the MuJoCo viewer only shows " + "collision shapes without it.", + ) + p.add_argument( + "--rebake", + action="store_true", + help="ignore cached wrapper.xml and OBJs; do a full re-bake.", + ) + p.add_argument( + "--view", + action="store_true", + help="launch the MuJoCo native viewer after baking (blocks, scene only — no robot).", + ) + args = p.parse_args() + + try: + scene_path = _resolve_existing_file(args.scene, "scene mesh") + except (FileNotFoundError, ValueError) as exc: + p.error(str(exc)) + + align = SceneMeshAlignment(scale=args.scale, y_up=not args.no_y_up) + spec = ( + CollisionSpec.from_json(args.collision_spec) + if args.collision_spec is not None + else CollisionSpec.auto_discover(scene_path) + ) + + wrapper = bake_scene_mjcf( + scene_mesh_path=scene_path, + alignment=align, + collision_spec=spec, + include_visual_mesh=args.visual, + rebake=args.rebake, + ) + print(f"wrapper: {wrapper}") + + if args.view: + import mujoco # type: ignore[import-untyped] + import mujoco.viewer # type: ignore[import-untyped] + + viewer: Any = mujoco.viewer + model = mujoco.MjModel.from_xml_path(str(wrapper)) + print(f"loaded: {model.nbody} bodies, {model.ngeom} geoms, {model.nmesh} meshes") + print("\nlaunching MuJoCo viewer (scene only — no robot attached)") + viewer.launch(model) + + +if __name__ == "__main__": + cli_main() + + +__all__ = ["bake_scene_mjcf", "load_or_bake"] diff --git a/dimos/simulation/mujoco/test_collision_spec.py b/dimos/simulation/mujoco/test_collision_spec.py new file mode 100644 index 0000000000..670f8a9037 --- /dev/null +++ b/dimos/simulation/mujoco/test_collision_spec.py @@ -0,0 +1,65 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import pytest + +from dimos.simulation.mujoco.collision_spec import CollisionSpec, decide_for_prim + + +def _flat_square_floor() -> tuple[np.ndarray, np.ndarray]: + vertices = np.asarray( + [ + [-1.0, -1.0, 0.0], + [1.0, -1.0, 0.0], + [1.0, 1.0, 0.0], + [-1.0, 1.0, 0.0], + ], + dtype=np.float64, + ) + triangles = np.asarray([[0, 1, 2], [0, 2, 3]], dtype=np.int32) + return vertices, triangles + + +def test_box_override_min_thickness_preserves_floor_top() -> None: + vertices, triangles = _flat_square_floor() + spec = CollisionSpec( + prim_overrides={ + "Floor*": { + "type": "box", + "min_thickness": 0.04, + "preserve": "top", + } + } + ) + + decision = decide_for_prim(vertices, triangles, "Floor_Plane.002", spec) + + assert decision.mode == "primitive" + assert decision.reason == "sidecar:box" + assert decision.primitive is not None + assert decision.primitive["size"] == pytest.approx((1.0, 1.0, 0.02)) + assert decision.primitive["pos"] == pytest.approx((0.0, 0.0, -0.02)) + + +def test_box_override_without_min_thickness_keeps_default_box_fit() -> None: + vertices, triangles = _flat_square_floor() + spec = CollisionSpec(prim_overrides={"Floor*": {"type": "box"}}) + + decision = decide_for_prim(vertices, triangles, "Floor_Plane.002", spec) + + assert decision.mode == "primitive" + assert decision.primitive is not None + assert min(decision.primitive["size"]) == pytest.approx(0.001) + assert decision.primitive["pos"][2] == pytest.approx(0.0) diff --git a/dimos/simulation/scene_assets/mesh_scene.py b/dimos/simulation/scene_assets/mesh_scene.py new file mode 100644 index 0000000000..ba16f7f396 --- /dev/null +++ b/dimos/simulation/scene_assets/mesh_scene.py @@ -0,0 +1,710 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Load a 3D scene asset into DimOS world-frame geometry. + +Supports: + * ``.glb`` / ``.gltf`` / ``.obj`` / ``.ply`` / ``.stl`` — via Open3D's + ``read_triangle_mesh``. + * ``.usdz`` / ``.usd`` / ``.usdc`` — via ``pxr.Usd`` (install ``usd-core``). + +Returned meshes are in DimOS world frame, with optional scale, +Y-up-to-Z-up rotation, Euler rotation, and translation applied. + +This loader is intentionally physics/viewer agnostic. MuJoCo collision +baking, browser collision baking, ray-casting, and asset inspection all +share the same source transform instead of each subsystem guessing its own +coordinate convention. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +import re +from typing import Any + +import numpy as np +import open3d as o3d # type: ignore[import-untyped] + +from dimos.simulation.scene_assets.spec import SceneMeshAlignment + +_TRIMESH_DUPLICATE_SUFFIX_RE = re.compile(r"_[0-9a-f]{6}$", re.IGNORECASE) + + +def _world_rotation(alignment: SceneMeshAlignment) -> np.ndarray: + """Compose the y-up swap + ZYX Euler into one 3x3.""" + rad = np.radians(alignment.rotation_zyx_deg) + cz, sz = np.cos(rad[0]), np.sin(rad[0]) + cy, sy = np.cos(rad[1]), np.sin(rad[1]) + cx, sx = np.cos(rad[2]), np.sin(rad[2]) + rz = np.array([[cz, -sz, 0], [sz, cz, 0], [0, 0, 1]], dtype=np.float64) + ry = np.array([[cy, 0, sy], [0, 1, 0], [-sy, 0, cy]], dtype=np.float64) + rx = np.array([[1, 0, 0], [0, cx, -sx], [0, sx, cx]], dtype=np.float64) + rzyx = rz @ ry @ rx + if alignment.y_up: + y_to_z = np.array( + [[1, 0, 0], [0, 0, -1], [0, 1, 0]], + dtype=np.float64, + ) + return rzyx @ y_to_z + return rzyx + + +def _average_per_face_vertex( + per_fv: np.ndarray, face_verts: np.ndarray, n_verts: int +) -> np.ndarray: + """Scatter-average ``(n_face_verts, 3)`` values onto ``(n_verts, 3)`` indices.""" + out = np.zeros((n_verts, 3), dtype=np.float32) + counts = np.zeros(n_verts, dtype=np.int32) + np.add.at(out, face_verts, per_fv) + np.add.at(counts, face_verts, 1) + counts = np.maximum(counts, 1)[:, None] + return out / counts + + +def _color_from_displaycolor( + mesh: Any, + n_verts: int, + face_counts: np.ndarray, + face_verts: np.ndarray, +) -> np.ndarray | None: + """Per-vertex RGB from ``primvars:displayColor`` if present and valued. + + Handles the four standard interpolations: ``constant`` / ``vertex`` / + ``uniform`` / ``faceVarying``. Returns ``None`` when the primvar + isn't authored with a value (Sketchfab USDZ exports typically declare + the primvar but leave it empty — colors live on the bound material). + """ + from pxr import UsdGeom # type: ignore[import-not-found, import-untyped] + + pv = UsdGeom.PrimvarsAPI(mesh.GetPrim()).GetPrimvar("displayColor") + if not pv or not pv.HasValue(): + return None + raw = pv.Get() + if raw is None: + return None + colors = np.asarray(raw, dtype=np.float32) + if colors.ndim != 2 or colors.shape[1] != 3 or colors.size == 0: + return None + interp = pv.GetInterpolation() + + if interp == UsdGeom.Tokens.constant: + return np.tile(colors[0:1], (n_verts, 1)) + + if interp == UsdGeom.Tokens.vertex and len(colors) == n_verts: + return colors + + if interp == UsdGeom.Tokens.uniform and len(colors) == len(face_counts): + per_fv = np.repeat(colors, face_counts, axis=0) + return _average_per_face_vertex(per_fv, face_verts, n_verts) + + if interp == UsdGeom.Tokens.faceVarying and len(colors) == len(face_verts): + return _average_per_face_vertex(colors, face_verts, n_verts) + + return None + + +def _color_from_material( + prim: Any, material_color_cache: dict[str, np.ndarray | None] +) -> np.ndarray | None: + """Per-prim RGB from the bound material's ``inputs:diffuseColor``. + + Walks ``UsdShadeMaterialBindingAPI`` → surface shader → ``inputs:diffuseColor``, + handling ``UsdPreviewSurface`` (the format Sketchfab USDZ uses). Texture + inputs aren't sampled — if ``diffuseColor`` is connected to a ``UsdUVTexture`` + rather than authored as a literal, this returns ``None`` and the caller + falls back to the next strategy. + + Results are cached per material path so we don't re-walk the shader graph + for every prim that shares a material. + """ + from pxr import UsdShade # type: ignore[import-not-found, import-untyped] + + mat_api = UsdShade.MaterialBindingAPI(prim) + bound = mat_api.ComputeBoundMaterial()[0] + if not bound: + return None + mat_path = str(bound.GetPath()) + if mat_path in material_color_cache: + return material_color_cache[mat_path] + + color = _resolve_diffuse_color(bound) + material_color_cache[mat_path] = color + return color + + +def _resolve_diffuse_color(material: Any) -> np.ndarray | None: + """Pull a literal ``diffuseColor`` out of a UsdShade material's surface shader.""" + from pxr import UsdShade # type: ignore[import-not-found, import-untyped] + + surface = material.ComputeSurfaceSource("")[0] + if not surface: + return None + diffuse_input = surface.GetInput("diffuseColor") + if not diffuse_input: + return None + # If the input is connected (texture-driven), bail — we don't sample images. + if diffuse_input.HasConnectedSource(): + connected = diffuse_input.GetConnectedSource()[0] + if connected: + shader = UsdShade.Shader(connected.GetPrim()) + if shader and shader.GetIdAttr().Get() == "UsdUVTexture": + return None + val = diffuse_input.Get() + if val is None: + return None + arr = np.asarray(val, dtype=np.float32).reshape(-1) + if arr.size != 3: + return None + return arr # (3,) RGB in [0, 1] + + +def _load_usd_mesh(path: Path) -> o3d.geometry.TriangleMesh: + """Walk every Mesh prim in a USD stage and concatenate to one o3d mesh. + + Also extracts per-vertex colors from ``primvars:displayColor`` when + present so downstream consumers can render textured-looking Sketchfab + exports without having to chase materials/textures. + """ + try: + from pxr import Usd, UsdGeom # type: ignore[import-not-found, import-untyped] + except ImportError as e: + raise ImportError("loading .usdz/.usd requires usd-core: `uv pip install usd-core`") from e + + stage = Usd.Stage.Open(str(path)) + if stage is None: + raise RuntimeError(f"could not open USD stage: {path}") + + all_pts: list[np.ndarray] = [] + all_tris: list[np.ndarray] = [] + all_colors: list[np.ndarray] = [] + any_color = False + vtx_offset = 0 + material_color_cache: dict[str, np.ndarray | None] = {} + + for prim in stage.Traverse(): + if not prim.IsA(UsdGeom.Mesh): + continue + mesh = UsdGeom.Mesh(prim) + pts_attr = mesh.GetPointsAttr().Get() + if pts_attr is None or len(pts_attr) == 0: + continue + pts = np.asarray(pts_attr, dtype=np.float32) + face_verts = np.asarray(mesh.GetFaceVertexIndicesAttr().Get(), dtype=np.int32) + face_counts = np.asarray(mesh.GetFaceVertexCountsAttr().Get(), dtype=np.int32) + + # Bake the prim's local-to-world transform into the points so the + # composite scene comes out in stage-root coordinates. + xform = UsdGeom.Xformable(prim).ComputeLocalToWorldTransform(Usd.TimeCode.Default()) + m = np.asarray(xform, dtype=np.float64).T # USD matrices are row-major + pts_h = np.hstack([pts, np.ones((len(pts), 1), dtype=np.float32)]) + pts_world = (m @ pts_h.T).T[:, :3].astype(np.float32) + + # Per-prim color resolution. Try in order: + # 1. ``primvars:displayColor`` (vertex / faceVarying / uniform / constant) + # 2. Bound material's ``inputs:diffuseColor`` (UsdPreviewSurface — what + # Sketchfab USDZ uses, with one constant color per material). + # 3. Neutral grey fallback. + prim_colors = _color_from_displaycolor(mesh, len(pts), face_counts, face_verts) + if prim_colors is None: + mat_color = _color_from_material(prim, material_color_cache) + if mat_color is not None: + prim_colors = np.tile(mat_color[None, :], (len(pts), 1)) + if prim_colors is not None: + any_color = True + else: + prim_colors = np.full((len(pts), 3), 0.7, dtype=np.float32) + + # USD allows quads / n-gons; fan-triangulate so o3d gets pure tris. + tris: list[tuple[int, int, int]] = [] + cursor = 0 + for n in face_counts: + for k in range(1, n - 1): + tris.append( + ( + int(face_verts[cursor]) + vtx_offset, + int(face_verts[cursor + k]) + vtx_offset, + int(face_verts[cursor + k + 1]) + vtx_offset, + ) + ) + cursor += n + + if not tris: + continue + all_pts.append(pts_world) + all_tris.append(np.asarray(tris, dtype=np.int32)) + all_colors.append(prim_colors) + vtx_offset += len(pts_world) + + if not all_pts: + raise RuntimeError(f"no Mesh prims with triangles found in {path}") + + pts = np.concatenate(all_pts, axis=0).astype(np.float64) + tris = np.concatenate(all_tris, axis=0) + + mesh = o3d.geometry.TriangleMesh() + mesh.vertices = o3d.utility.Vector3dVector(pts) + mesh.triangles = o3d.utility.Vector3iVector(tris) + if any_color: + colors = np.concatenate(all_colors, axis=0).astype(np.float64) + mesh.vertex_colors = o3d.utility.Vector3dVector(np.clip(colors, 0.0, 1.0)) + return mesh + + +def load_scene_mesh( + path: str | Path, + alignment: SceneMeshAlignment | None = None, +) -> o3d.geometry.TriangleMesh: + """Load a scene mesh from disk and apply alignment to put it in dimos world frame. + + Args: + path: file path. Supported extensions: ``.usdz``, ``.usd``, ``.usdc``, + ``.glb``, ``.gltf``, ``.obj``, ``.ply``, ``.stl``. + alignment: scale / rotation / translation to apply. + + Returns: + an ``open3d.geometry.TriangleMesh`` in dimos world frame with vertex + normals computed. + """ + path = Path(path) + if not path.exists(): + raise FileNotFoundError(f"scene mesh not found: {path}") + suffix = path.suffix.lower() + if suffix in {".usdz", ".usd", ".usdc", ".usda"}: + mesh = _load_usd_mesh(path) + elif suffix in {".glb", ".gltf"}: + # GEOMETRY-ONLY GLB load. Used by floor-z probing and ray-casting; + # it does not need PBR materials. ``trimesh.load(path, force="mesh")`` + # would flatten the scene by decompressing every embedded texture and + # sampling per-vertex colors. For a scene with hundreds of 4K PBR + # textures, that allocates ~10 GB transiently and OOMs 32 GB boxes. + # We open in Scene mode (no flattening, no texture decode), walk the + # instance graph applying each instance's world transform, and emit a + # single concatenated mesh — peak stays under ~1 GB. + import trimesh + + scene_or_mesh: Any = trimesh.load(str(path)) + if isinstance(scene_or_mesh, trimesh.Trimesh): + verts_world = np.asarray(scene_or_mesh.vertices, dtype=np.float64) + faces_world = np.asarray(scene_or_mesh.faces, dtype=np.int64) + else: + scene = scene_or_mesh + verts_chunks: list[np.ndarray] = [] + faces_chunks: list[np.ndarray] = [] + v_off = 0 + for node_name in scene.graph.nodes_geometry: + xform, geom_name = scene.graph[node_name] + geom = scene.geometry.get(geom_name) + if geom is None or not isinstance(geom, trimesh.Trimesh) or len(geom.faces) == 0: + continue + v_local = np.asarray(geom.vertices, dtype=np.float64) + f_local = np.asarray(geom.faces, dtype=np.int64) + m = np.asarray(xform, dtype=np.float64) + v_h = np.hstack([v_local, np.ones((len(v_local), 1), dtype=np.float64)]) + v_world = (m @ v_h.T).T[:, :3] + verts_chunks.append(v_world) + faces_chunks.append(f_local + v_off) + v_off += len(v_local) + if not verts_chunks: + raise RuntimeError(f"glTF loaded but no Trimesh instances found: {path}") + verts_world = np.concatenate(verts_chunks, axis=0) + faces_world = np.concatenate(faces_chunks, axis=0) + + mesh = o3d.geometry.TriangleMesh() + mesh.vertices = o3d.utility.Vector3dVector(verts_world) + mesh.triangles = o3d.utility.Vector3iVector(faces_world.astype(np.int32)) + else: + mesh = o3d.io.read_triangle_mesh(str(path)) + if len(mesh.triangles) == 0: + raise RuntimeError(f"o3d.io.read_triangle_mesh returned an empty mesh for {path}") + + align = alignment or SceneMeshAlignment() + if align.scale != 1.0: + mesh.scale(align.scale, center=np.zeros(3)) + rot = _world_rotation(align) + if not np.allclose(rot, np.eye(3)): + mesh.rotate(rot, center=np.zeros(3)) + if any(align.translation): + mesh.translate(np.asarray(align.translation, dtype=np.float64)) + + mesh.compute_vertex_normals() + return mesh + + +def floor_z_under_origin( + scene_mesh_path: str | Path, + alignment: SceneMeshAlignment | None = None, +) -> float: + """Return the first scene surface under world ``x=0, y=0``. + + Falls back to the mesh bbox minimum when the origin ray misses. + """ + import open3d.core as o3c # type: ignore[import-untyped] + + mesh = load_scene_mesh(scene_mesh_path, alignment=alignment) + scene = make_raycasting_scene(mesh) + rays = o3c.Tensor( + np.array([[0.0, 0.0, 1000.0, 0.0, 0.0, -1.0]], dtype=np.float32), + dtype=o3c.Dtype.Float32, + ) + t_hit = float(scene.cast_rays(rays)["t_hit"].numpy()[0]) + if np.isfinite(t_hit): + return 1000.0 - t_hit + bbox = mesh.get_axis_aligned_bounding_box() + return float(bbox.min_bound[2]) + + +def make_raycasting_scene( + mesh: o3d.geometry.TriangleMesh, +) -> o3d.t.geometry.RaycastingScene: + """Wrap a TriangleMesh into Open3D's BVH-backed ray-casting scene.""" + scene = o3d.t.geometry.RaycastingScene() + scene.add_triangles(o3d.t.geometry.TriangleMesh.from_legacy(mesh)) + return scene + + +@dataclass +class ScenePrimMesh: + """One USD ``Mesh`` prim's geometry, ready to write to OBJ. + + Used by ``load_scene_prims`` to keep prims separate so MuJoCo can + treat each as its own (approximately convex) collision shape. When + the loader handles a non-USD format the input is returned as a + single-element list with the whole mesh in it. + """ + + name: str + """Sanitized identifier (safe for MJCF asset names) — typically the + USD prim path with non-alphanumerics replaced.""" + + vertices: np.ndarray + """``(N, 3)`` float32, in world frame after alignment.""" + + triangles: np.ndarray + """``(M, 3)`` int32 vertex indices.""" + + prim_path: str | None = None + """Original scene-graph path when the source format provides one.""" + + visual_node_name: str | None = None + """Stable source node name used by visual extraction when available.""" + + +def split_disconnected_scene_prims( + prims: list[ScenePrimMesh], + *, + min_components: int, + extent_ratio: float, + prim_min_extent: float, + axis_ratio: float, + min_component_extent: float, + min_component_faces: int, + can_split: Callable[[ScenePrimMesh], bool] | None = None, +) -> tuple[list[ScenePrimMesh], dict[str, int]]: + """Split scene-graph nodes that are disconnected prop clusters. + + Some game exports group many small disconnected objects under one node + (for example leaves, cups, bottles). Primitive fitting sees only the + combined bounds and can turn the group into one scene-scale slab. This + helper keeps normal connected props intact, but splits suspicious wide + clusters so tiny decorative islands can be dropped and larger islands can + be fitted independently. + """ + import trimesh + + result: list[ScenePrimMesh] = [] + stats = { + "source_prims": len(prims), + "split_prims": 0, + "emitted_components": 0, + "dropped_components": 0, + } + + for prim in prims: + if can_split is not None and not can_split(prim): + result.append(prim) + continue + if len(prim.triangles) < max(min_component_faces * 2, 1): + result.append(prim) + continue + prim_extent = np.ptp(prim.vertices, axis=0) + if float(prim_extent.max()) < prim_min_extent: + result.append(prim) + continue + positive_extent = prim_extent[prim_extent > 1e-6] + if ( + len(positive_extent) < 3 + or float(positive_extent.max() / positive_extent.min()) < axis_ratio + ): + result.append(prim) + continue + + mesh = trimesh.Trimesh(vertices=prim.vertices, faces=prim.triangles, process=False) + parts = mesh.split(only_watertight=False) + if len(parts) < min_components: + result.append(prim) + continue + + component_extents = np.array( + [np.ptp(np.asarray(part.vertices), axis=0).max() for part in parts], + dtype=np.float64, + ) + median_component_extent = float(np.median(component_extents)) + if median_component_extent <= 0.0: + result.append(prim) + continue + if float(prim_extent.max()) / median_component_extent < extent_ratio: + result.append(prim) + continue + + emitted = 0 + dropped = 0 + for index, part in enumerate(parts): + vertices = np.asarray(part.vertices, dtype=np.float32) + triangles = np.asarray(part.faces, dtype=np.int32) + component_extent = float(np.ptp(vertices, axis=0).max()) if len(vertices) else 0.0 + if len(triangles) < min_component_faces or component_extent < min_component_extent: + dropped += 1 + continue + result.append( + ScenePrimMesh( + name=f"{prim.name}_part{index:04d}", + vertices=vertices, + triangles=triangles, + prim_path=( + f"{prim.prim_path}/component_{index:04d}" + if prim.prim_path is not None + else f"{prim.name}/component_{index:04d}" + ), + ) + ) + emitted += 1 + + stats["split_prims"] += 1 + stats["emitted_components"] += emitted + stats["dropped_components"] += dropped + if emitted == 0: + continue + + return result, stats + + +def _load_glb_prims(path: Path, alignment: SceneMeshAlignment) -> list[ScenePrimMesh]: + """Enumerate per-instance prims from a glTF/GLB. + + ``trimesh.load(file.glb)`` returns a ``Scene`` whose ``graph`` records + the world transform for every geometry instance. Iterating + ``graph.nodes_geometry`` is the trimesh equivalent of USD's + ``stage.Traverse()`` — it yields one entry per instance, even when + multiple instances share the same underlying mesh (typical for chairs, + cabinets, etc.). Without this enumeration, ``trimesh.load(... force="mesh")`` + collapses the whole scene to one mesh and CoACD produces a single coarse + decomposition, which is essentially useless for collision against + multi-object scenes. + """ + import trimesh + + loaded: Any = trimesh.load(str(path)) + R = _world_rotation(alignment) + T = np.asarray(alignment.translation, dtype=np.float64) + s = float(alignment.scale) + + if isinstance(loaded, trimesh.Trimesh): + # Single-mesh GLB (no scene graph). Treat as one prim. + pts = np.asarray(loaded.vertices, dtype=np.float64) + faces = np.asarray(loaded.faces, dtype=np.int32) + if len(faces) == 0: + return [] + pts_world = (R @ (s * pts).T).T + T + return [ + ScenePrimMesh( + name="scene", + vertices=pts_world.astype(np.float32), + triangles=faces, + prim_path="scene", + ) + ] + + scene = loaded + prims: list[ScenePrimMesh] = [] + name_counts: dict[str, int] = {} + prim_path_counts: dict[str, int] = {} + for node_name in scene.graph.nodes_geometry: + xform, geom_name = scene.graph[node_name] + geom = scene.geometry.get(geom_name) + if geom is None or not isinstance(geom, trimesh.Trimesh): + continue + if len(geom.faces) == 0: + continue + + pts_local = np.asarray(geom.vertices, dtype=np.float64) + faces = np.asarray(geom.faces, dtype=np.int32) + + # Local → scene-root via the instance transform. + m = np.asarray(xform, dtype=np.float64) + pts_h = np.hstack([pts_local, np.ones((len(pts_local), 1), dtype=np.float64)]) + pts_stage = (m @ pts_h.T).T[:, :3] + + # Scene-root → dimos world via SceneMeshAlignment. + pts_world = (R @ (s * pts_stage).T).T + T + + stable_node = _stable_trimesh_node_name(str(node_name)) + stable_prim_path = _unique_stable_name( + f"{stable_node}_{geom_name}", + prim_path_counts, + ) + clean = _unique_stable_name(_sanitize_scene_name(stable_prim_path), name_counts) + prims.append( + ScenePrimMesh( + name=clean, + vertices=pts_world.astype(np.float32), + triangles=faces, + prim_path=stable_prim_path, + visual_node_name=stable_node, + ) + ) + return sorted(prims, key=lambda prim: prim.prim_path or prim.name) + + +def _stable_trimesh_node_name(node_name: str) -> str: + """Drop random duplicate suffixes that trimesh adds to glTF nodes.""" + return _TRIMESH_DUPLICATE_SUFFIX_RE.sub("", node_name) + + +def _sanitize_scene_name(raw: str) -> str: + return "".join(c if c.isalnum() else "_" for c in raw) + + +def _unique_stable_name(raw: str, counts: dict[str, int]) -> str: + count = counts.get(raw, 0) + counts[raw] = count + 1 + if count == 0: + return raw + return f"{raw}__{count:03d}" + + +def load_scene_prims( + path: str | Path, + alignment: SceneMeshAlignment | None = None, +) -> list[ScenePrimMesh]: + """Load a USD/USDZ scene as one ``ScenePrimMesh`` per Mesh prim. + + Per-prim splitting is what MuJoCo wants for non-trivial scenes: + each prim's convex hull approximates the prim well, while the + convex hull of the *whole* scene is its bounding box. Falls back + to a single ScenePrimMesh for non-USD inputs (a single ``.obj`` or + ``.glb`` doesn't carry per-part semantics in our loader). + + Same alignment rules as ``load_scene_mesh``. + """ + path = Path(path) + align = alignment or SceneMeshAlignment() + suffix = path.suffix.lower() + + if suffix in {".glb", ".gltf"}: + return _load_glb_prims(path, align) + + if suffix not in {".usdz", ".usd", ".usdc", ".usda"}: + # Non-USD, non-glTF (e.g. .obj/.ply/.stl): one part, whole mesh. + whole = load_scene_mesh(path, alignment=align) + return [ + ScenePrimMesh( + name="scene", + vertices=np.asarray(whole.vertices, dtype=np.float32), + triangles=np.asarray(whole.triangles, dtype=np.int32), + prim_path="scene", + ) + ] + + try: + from pxr import Usd, UsdGeom # type: ignore[import-not-found, import-untyped] + except ImportError as e: + raise ImportError("loading .usdz/.usd requires usd-core: `uv pip install usd-core`") from e + + stage = Usd.Stage.Open(str(path)) + if stage is None: + raise RuntimeError(f"could not open USD stage: {path}") + + R = _world_rotation(align) + T = np.asarray(align.translation, dtype=np.float64) + s = float(align.scale) + + prims: list[ScenePrimMesh] = [] + for prim in stage.Traverse(): + if not prim.IsA(UsdGeom.Mesh): + continue + usd_mesh = UsdGeom.Mesh(prim) + pts_attr = usd_mesh.GetPointsAttr().Get() + if pts_attr is None or len(pts_attr) == 0: + continue + pts = np.asarray(pts_attr, dtype=np.float64) + face_verts = np.asarray(usd_mesh.GetFaceVertexIndicesAttr().Get(), dtype=np.int32) + face_counts = np.asarray(usd_mesh.GetFaceVertexCountsAttr().Get(), dtype=np.int32) + + # Local → stage-root via the USD prim's accumulated transform. + xform = UsdGeom.Xformable(prim).ComputeLocalToWorldTransform(Usd.TimeCode.Default()) + m = np.asarray(xform, dtype=np.float64).T + pts_h = np.hstack([pts, np.ones((len(pts), 1), dtype=np.float64)]) + pts_stage = (m @ pts_h.T).T[:, :3] + + # Stage-root → dimos world via SceneMeshAlignment (scale → rot → trans). + pts_world = (R @ (s * pts_stage).T).T + T + + # Triangulate any quads / n-gons (vertex indices are local to this prim now). + tris: list[tuple[int, int, int]] = [] + cursor = 0 + for n in face_counts: + for k in range(1, n - 1): + tris.append( + ( + int(face_verts[cursor]), + int(face_verts[cursor + k]), + int(face_verts[cursor + k + 1]), + ) + ) + cursor += n + if not tris: + continue + + # MJCF asset names: strip the leading slash, swap remaining + # path separators / dots for underscores. USD prim paths can + # collide on the same leaf; suffix the index so each is unique. + raw = str(prim.GetPath()).lstrip("/") + clean = "".join(c if c.isalnum() else "_" for c in raw) + prim_path = str(prim.GetPath()) + prims.append( + ScenePrimMesh( + name=f"{clean}__{len(prims)}", + vertices=pts_world.astype(np.float32), + triangles=np.asarray(tris, dtype=np.int32), + prim_path=prim_path, + ) + ) + + if not prims: + raise RuntimeError(f"no Mesh prims with triangles found in {path}") + return prims + + +__all__ = [ + "SceneMeshAlignment", + "ScenePrimMesh", + "floor_z_under_origin", + "load_scene_mesh", + "load_scene_prims", + "make_raycasting_scene", + "split_disconnected_scene_prims", +] diff --git a/pyproject.toml b/pyproject.toml index e4f03e5406..86be7568f5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -313,8 +313,18 @@ apriltag = [ "reportlab>=4.5.0", ] +# Offline scene-package cooking + loading: mesh/USD import (open3d, usd-core), +# convex decomposition for collision (coacd), and mesh ops (trimesh). The +# Blender visual bake is an optional external `blender` binary, not a wheel. +scene = [ + "open3d>=0.18.0", + "trimesh>=4.0.0", + "coacd>=1.0.0", + "usd-core>=23.11", +] + all = [ - "dimos[agents,apriltag,base,cpu,cuda,drone,manipulation,misc,perception,sim,unitree,visualization,web]", + "dimos[agents,apriltag,base,cpu,cuda,drone,manipulation,misc,perception,scene,sim,unitree,visualization,web]", ] [dependency-groups] @@ -548,6 +558,13 @@ module = [ "websocket", "xarm.*", "ament_index_python.*", + "coacd", + "open3d", + "open3d.*", + "pxr", + "pxr.*", + "trimesh", + "trimesh.*", ] ignore_missing_imports = true diff --git a/uv.lock b/uv.lock index a22bbafc62..10f8a26e12 100644 --- a/uv.lock +++ b/uv.lock @@ -1256,6 +1256,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/56/f3/4da9d5c5308ef2019ab65a8a9f519ac95004446902d01e859f9ac6b8cdd6/cmeel_zlib-1.3.1-0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:1e36ac8dccca22ff1f6e4df428ae5597f6288d9e6f85b08c9b767dc63e90fb55", size = 285662, upload-time = "2025-02-11T12:20:37.298Z" }, ] +[[package]] +name = "coacd" +version = "1.0.11" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/1d/3b/434beab9e1754ca0ae3a619b383243dd85aa65d03a4dc7333c8296c97a92/coacd-1.0.11-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:adbc58259a721cc5ede24cc8c2671a95e75a8a52dc1ee4d953d80d236b192da9", size = 3299264, upload-time = "2026-05-04T19:23:36.183Z" }, + { url = "https://files.pythonhosted.org/packages/cb/a7/06f63baa29198f681ba60848e3271cc625eddd61cd4e59071f56ffce9362/coacd-1.0.11-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e376fadb22790444c7253f0cee9104a1af01ec965488c1318e84e3b2dbf1e2a3", size = 2573832, upload-time = "2026-05-04T19:23:38.008Z" }, + { url = "https://files.pythonhosted.org/packages/0e/b4/057c78f7b16b87871cfcecc6febe70522e77a2a33f8788960377152c224d/coacd-1.0.11-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a60f700a52e5b40462e508c14bb756cd63ce7e6a95ff72ae0b1592be1dbb0106", size = 2640170, upload-time = "2026-05-04T19:23:39.322Z" }, + { url = "https://files.pythonhosted.org/packages/00/83/c50472ce98175fddd86d4aba861fea89f30350a5487880b3a81d34915a85/coacd-1.0.11-cp39-abi3-win_amd64.whl", hash = "sha256:4de22f70d1a3fa8c44698c8006a223fe5fb0ee84b76adecf3726cf2003e9145f", size = 1465079, upload-time = "2026-05-04T19:23:41.604Z" }, +] + [[package]] name = "coal" version = "3.0.2" @@ -2015,6 +2030,7 @@ agents = [ all = [ { name = "a750-control", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "chromadb" }, + { name = "coacd" }, { name = "cupy-cuda12x", marker = "platform_machine == 'x86_64'" }, { name = "dimos-viewer" }, { name = "drake", version = "1.45.0", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine != 'aarch64' and sys_platform == 'darwin'" }, @@ -2044,6 +2060,7 @@ all = [ { name = "onnxruntime" }, { name = "onnxruntime-gpu", marker = "platform_machine == 'x86_64'" }, { name = "open-clip-torch" }, + { name = "open3d" }, { name = "openai" }, { name = "opencv-contrib-python" }, { name = "pillow" }, @@ -2070,6 +2087,7 @@ all = [ { name = "trimesh" }, { name = "ultralytics" }, { name = "unitree-webrtc-connect" }, + { name = "usd-core" }, { name = "uvicorn" }, { name = "viser", extra = ["urdf"] }, { name = "xacro" }, @@ -2163,6 +2181,12 @@ perception = [ { name = "transformers", extra = ["torch"] }, { name = "ultralytics" }, ] +scene = [ + { name = "coacd" }, + { name = "open3d" }, + { name = "trimesh" }, + { name = "usd-core" }, +] sim = [ { name = "mujoco" }, { name = "playground" }, @@ -2405,12 +2429,13 @@ requires-dist = [ { name = "a750-control", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'manipulation'" }, { name = "annotation-protocol", specifier = ">=1.4.0" }, { name = "bleak", specifier = ">=3.0.2" }, + { name = "coacd", marker = "extra == 'scene'", specifier = ">=1.0.0" }, { name = "chromadb", marker = "extra == 'perception'", specifier = ">=1.0.0" }, { name = "cryptography", specifier = ">=46.0.5" }, { name = "cupy-cuda12x", marker = "platform_machine == 'x86_64' and extra == 'cuda'", specifier = "==13.6.0" }, { name = "cyclonedds", marker = "extra == 'dds'", specifier = ">=0.10.5" }, { name = "cyclonedds", marker = "extra == 'unitree-dds'", specifier = ">=0.10.5" }, - { name = "dimos", extras = ["agents", "apriltag", "base", "cpu", "cuda", "drone", "manipulation", "misc", "perception", "sim", "unitree", "visualization", "web"], marker = "extra == 'all'" }, + { name = "dimos", extras = ["agents", "apriltag", "base", "cpu", "cuda", "drone", "manipulation", "misc", "perception", "scene", "sim", "unitree", "visualization", "web"], marker = "extra == 'all'" }, { name = "dimos", extras = ["agents", "web", "perception", "visualization"], marker = "extra == 'base'" }, { name = "dimos", extras = ["base", "mapping"], marker = "extra == 'unitree'" }, { name = "dimos", extras = ["unitree"], marker = "extra == 'unitree-dds'" }, @@ -2450,6 +2475,7 @@ requires-dist = [ { name = "onnxruntime-gpu", marker = "platform_machine == 'x86_64' and extra == 'cuda'", specifier = ">=1.17.1" }, { name = "open-clip-torch", marker = "extra == 'misc'", specifier = "==3.2.0" }, { name = "open3d", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'", specifier = ">=0.18.0" }, + { name = "open3d", marker = "extra == 'scene'", specifier = ">=0.18.0" }, { name = "open3d-unofficial-arm", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'", specifier = ">=0.19.0.post9" }, { name = "openai", marker = "extra == 'agents'" }, { name = "opencv-contrib-python", marker = "extra == 'apriltag'", specifier = "==4.10.0.84" }, @@ -2495,11 +2521,13 @@ requires-dist = [ { name = "torchreid", marker = "extra == 'misc'", specifier = "==0.2.5" }, { name = "transformers", extras = ["torch"], marker = "extra == 'perception'", specifier = ">=4.53.0,<4.54" }, { name = "trimesh", marker = "extra == 'manipulation'" }, + { name = "trimesh", marker = "extra == 'scene'", specifier = ">=4.0.0" }, { name = "typer", specifier = ">=0.19.2,<1" }, { name = "typing-extensions", marker = "python_full_version < '3.11'", specifier = ">=4.0" }, { name = "ultralytics", marker = "extra == 'perception'", specifier = ">=8.3.70" }, { name = "unitree-sdk2py-dimos", marker = "extra == 'unitree-dds'", specifier = ">=1.0.2" }, { name = "unitree-webrtc-connect", marker = "extra == 'unitree'", specifier = ">=2.1.2" }, + { name = "usd-core", marker = "extra == 'scene'", specifier = ">=23.11" }, { name = "uvicorn", marker = "extra == 'web'", specifier = ">=0.34.0" }, { name = "viser", extras = ["urdf"], marker = "extra == 'manipulation'", specifier = ">=1.0.29" }, { name = "websocket-client", specifier = ">=1.8" }, @@ -2507,7 +2535,7 @@ requires-dist = [ { name = "xarm-python-sdk", marker = "extra == 'manipulation'", specifier = ">=1.17.0" }, { name = "xarm-python-sdk", marker = "extra == 'misc'", specifier = ">=1.17.0" }, ] -provides-extras = ["misc", "visualization", "agents", "web", "perception", "unitree", "unitree-dds", "manipulation", "cpu", "cuda", "sim", "mapping", "drone", "dds", "base", "apriltag", "all"] +provides-extras = ["misc", "visualization", "agents", "web", "perception", "unitree", "unitree-dds", "manipulation", "cpu", "cuda", "sim", "mapping", "drone", "dds", "base", "apriltag", "scene", "all"] [package.metadata.requires-dev] autofix = [{ name = "ruff", specifier = "==0.14.3" }] @@ -6860,23 +6888,23 @@ name = "open3d" version = "0.19.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "addict", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" }, - { name = "configargparse", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" }, - { name = "dash", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" }, - { name = "flask", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" }, - { name = "matplotlib", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" }, - { name = "nbformat", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'aarch64') or (python_full_version < '3.11' and sys_platform != 'linux')" }, - { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'aarch64') or (python_full_version >= '3.11' and sys_platform != 'linux')" }, - { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'aarch64') or (python_full_version < '3.11' and sys_platform != 'linux')" }, - { name = "pandas", version = "3.0.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'aarch64') or (python_full_version >= '3.11' and sys_platform != 'linux')" }, - { name = "pillow", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" }, - { name = "pyquaternion", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" }, - { name = "pyyaml", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" }, - { name = "scikit-learn", version = "1.7.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'aarch64') or (python_full_version < '3.11' and sys_platform != 'linux')" }, - { name = "scikit-learn", version = "1.8.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'aarch64') or (python_full_version >= '3.11' and sys_platform != 'linux')" }, - { name = "tqdm", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" }, - { name = "werkzeug", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" }, + { name = "addict" }, + { name = "configargparse" }, + { name = "dash" }, + { name = "flask" }, + { name = "matplotlib" }, + { name = "nbformat" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "pandas", version = "3.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "pillow" }, + { name = "pyquaternion" }, + { name = "pyyaml" }, + { name = "scikit-learn", version = "1.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scikit-learn", version = "1.8.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "tqdm" }, + { name = "werkzeug" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/5c/4b/91e8a4100adf0ccd2f7ad21dd24c2e3d8f12925396528d0462cfb1735e5a/open3d-0.19.0-cp310-cp310-macosx_10_15_universal2.whl", hash = "sha256:f7128ded206e07987cc29d0917195fb64033dea31e0d60dead3629b33d3c175f", size = 103086005, upload-time = "2025-01-08T07:25:56.755Z" }, @@ -7305,16 +7333,17 @@ source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version < '3.11' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version < '3.11' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", "python_full_version < '3.11' and platform_machine == 'x86_64' and sys_platform == 'win32'", "python_full_version < '3.11' and platform_machine != 'x86_64' and sys_platform == 'win32'", "python_full_version < '3.11' and platform_machine == 'x86_64' and sys_platform != 'darwin' and sys_platform != 'win32'", "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (python_full_version < '3.11' and platform_machine != 'x86_64' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')", ] dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'aarch64') or (python_full_version < '3.11' and sys_platform != 'linux')" }, - { name = "python-dateutil", marker = "(python_full_version < '3.11' and platform_machine != 'aarch64') or (python_full_version < '3.11' and sys_platform != 'linux')" }, - { name = "pytz", marker = "(python_full_version < '3.11' and platform_machine != 'aarch64') or (python_full_version < '3.11' and sys_platform != 'linux')" }, - { name = "tzdata", marker = "(python_full_version < '3.11' and platform_machine != 'aarch64') or (python_full_version < '3.11' and sys_platform != 'linux')" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "python-dateutil", marker = "python_full_version < '3.11'" }, + { name = "pytz", marker = "python_full_version < '3.11'" }, + { name = "tzdata", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/33/01/d40b85317f86cf08d853a4f495195c73815fdf205eef3993821720274518/pandas-2.3.3.tar.gz", hash = "sha256:e05e1af93b977f7eafa636d043f9f94c7ee3ac81af99c13508215942e64c993b", size = 4495223, upload-time = "2025-09-29T23:34:51.853Z" } wheels = [ @@ -7378,6 +7407,9 @@ resolution-markers = [ "python_full_version == '3.13.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version >= '3.14' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.13.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'win32'", "python_full_version >= '3.14' and platform_machine != 'x86_64' and sys_platform == 'win32'", "python_full_version == '3.13.*' and platform_machine == 'x86_64' and sys_platform == 'win32'", @@ -7392,14 +7424,15 @@ resolution-markers = [ "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (python_full_version == '3.12.*' and platform_machine != 'x86_64' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')", "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'darwin'", "python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'win32'", "python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform == 'win32'", "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform != 'darwin' and sys_platform != 'win32'", "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (python_full_version == '3.11.*' and platform_machine != 'x86_64' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')", ] dependencies = [ - { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'aarch64') or (python_full_version >= '3.11' and sys_platform != 'linux')" }, - { name = "python-dateutil", marker = "(python_full_version >= '3.11' and platform_machine != 'aarch64') or (python_full_version >= '3.11' and sys_platform != 'linux')" }, + { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "python-dateutil", marker = "python_full_version >= '3.11'" }, { name = "tzdata", marker = "(python_full_version >= '3.11' and sys_platform == 'emscripten') or (python_full_version >= '3.11' and sys_platform == 'win32')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/de/da/b1dc0481ab8d55d0f46e343cfe67d4551a0e14fcee52bd38ca1bd73258d8/pandas-3.0.0.tar.gz", hash = "sha256:0facf7e87d38f721f0af46fe70d97373a37701b1c09f7ed7aeeb292ade5c050f", size = 4633005, upload-time = "2026-01-21T15:52:04.726Z" } @@ -8859,8 +8892,8 @@ name = "pyquaternion" version = "0.9.9" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'aarch64') or (python_full_version < '3.11' and sys_platform != 'linux')" }, - { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'aarch64') or (python_full_version >= '3.11' and sys_platform != 'linux')" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7d/0d/3d092aa20efaedacb89c3221a92c6491be5b28f618a2c36b52b53e7446c2/pyquaternion-0.9.9.tar.gz", hash = "sha256:b1f61af219cb2fe966b5fb79a192124f2e63a3f7a777ac3cadf2957b1a81bea8", size = 15530, upload-time = "2020-10-05T01:31:30.327Z" } wheels = [ @@ -11403,6 +11436,28 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl", hash = "sha256:bf272323e553dfb2e87d9bfd225ca7b0f467b919d7bbd355436d3fd37cb0acd4", size = 131584, upload-time = "2026-01-07T16:24:42.685Z" }, ] +[[package]] +name = "usd-core" +version = "26.5" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9b/a0/639e148c16a0ec201cc4848aa3da4aba8805e17a2d9e2398eec399fd3051/usd_core-26.5-cp310-none-macosx_10_15_universal2.whl", hash = "sha256:d6a3a567e313841b7390ea7a930bf5aef08bdb912974c725becd725d83edb0f9", size = 39723088, upload-time = "2026-04-24T20:17:23.663Z" }, + { url = "https://files.pythonhosted.org/packages/d7/26/6cb620a64f3fafa38b84008d916eee47c70e5313c5d88c9087edf4d57522/usd_core-26.5-cp310-none-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:85a1484024cdcefd77aac32a3b98e698f655e01951d62cc4d3fb3826e232400c", size = 28820064, upload-time = "2026-04-24T20:17:27.161Z" }, + { url = "https://files.pythonhosted.org/packages/00/d7/7814c95ca0b13a26313e5256472f90cfa2ab7f7cf3103b0d3611d41156e6/usd_core-26.5-cp310-none-win_amd64.whl", hash = "sha256:dff985cbfe24870a5dfe1c578acd918a358cd1680a17777d83b55d50f5560c18", size = 13450099, upload-time = "2026-04-24T20:17:29.994Z" }, + { url = "https://files.pythonhosted.org/packages/39/3a/adf7a4043e70974b84d3a572f928ffdd1176a070595cd17f028062622ade/usd_core-26.5-cp311-none-macosx_10_15_universal2.whl", hash = "sha256:b5416a108080311632b975da71b4ea480757ac6e7ea19b30bcd0eed6a3b6081f", size = 39723550, upload-time = "2026-04-24T20:17:32.975Z" }, + { url = "https://files.pythonhosted.org/packages/e2/7f/575b0ddc2a3effa1dc1f50ed67ae0def8f9ed961c69bfbb89a0a1c9ceaf8/usd_core-26.5-cp311-none-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:60076c97f0de2611dc39d2d25826e3b22a2b0e391c73806b4a072d69929f329e", size = 28825210, upload-time = "2026-04-24T20:17:37.136Z" }, + { url = "https://files.pythonhosted.org/packages/9f/51/9fb7c817f1ee7aff02adde8ec4805ff4add06482e036fe0914ab8e9cdbc5/usd_core-26.5-cp311-none-win_amd64.whl", hash = "sha256:1ff2031095ecdc2f9ff4e245114e6ab7001f7dec8fe75436b5beb72e1a280f57", size = 13450734, upload-time = "2026-04-24T20:17:39.641Z" }, + { url = "https://files.pythonhosted.org/packages/8d/cc/04870cc3ae8e1b3a4e168efea47e389cfab6ab4f619005da2443a10390d4/usd_core-26.5-cp312-none-macosx_10_15_universal2.whl", hash = "sha256:a9df2864e84b83ffc9cc0f2777a49170180f84f2b679bcd014d72036a51d057c", size = 39775789, upload-time = "2026-04-24T20:17:43.025Z" }, + { url = "https://files.pythonhosted.org/packages/77/62/963d3aba966539917d01e4a2169a1c07f7b3df087fc16ee39fc764214969/usd_core-26.5-cp312-none-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:caa2447252aeada8c158faacd4d448f29cf1617aeccef5bb954734b93c8f3f62", size = 28743527, upload-time = "2026-04-24T20:17:46.631Z" }, + { url = "https://files.pythonhosted.org/packages/3b/b0/645ae6e27a9768e570c1044efd6d2369c10c5c2412669314b3d6cd914803/usd_core-26.5-cp312-none-win_amd64.whl", hash = "sha256:6d887b010c756508d2e1f770626201f1f4ba5227c052c1135ba9c19932c4da8e", size = 13494028, upload-time = "2026-04-24T20:17:49.599Z" }, + { url = "https://files.pythonhosted.org/packages/2e/cd/128de2e16d597eb0868dde7cc837a908b28ec2a0d90d4697714b6770449b/usd_core-26.5-cp313-none-macosx_10_15_universal2.whl", hash = "sha256:ce5e90a6795b93d7e744694e5209ea2f1754f9d596e67a89f0cc3590e9fff578", size = 39776038, upload-time = "2026-04-24T20:17:52.535Z" }, + { url = "https://files.pythonhosted.org/packages/f1/10/88838fd371592cfc3d972547ab4361e2deef5891d89c22a509de0e6696ce/usd_core-26.5-cp313-none-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:dd4d3de388e6dfec91fa5ee9fa29800d43ebe86cbf7a10380ec02b15386fca67", size = 28743992, upload-time = "2026-04-24T20:17:55.995Z" }, + { url = "https://files.pythonhosted.org/packages/62/05/da8f44024e0f947c13da3bdae0d4ac6c04cb86de92a6f1b9bf03e6bb8ae8/usd_core-26.5-cp313-none-win_amd64.whl", hash = "sha256:b077ea37dfeb15ca6b24ca33b65c2fe9b1656138e1fda74e4eae9793a149a7d5", size = 13494201, upload-time = "2026-04-24T20:17:59.015Z" }, + { url = "https://files.pythonhosted.org/packages/3d/57/01cc4e412feaad5aaee09d09ead2afbd5b4022e3d3b5461adcbf726ca3f8/usd_core-26.5-cp314-none-macosx_10_15_universal2.whl", hash = "sha256:5b0acd9a1d804cb73d58815365ccb141727f635f4e6764609fade3bf4ef5cbba", size = 39927684, upload-time = "2026-04-24T20:18:01.828Z" }, + { url = "https://files.pythonhosted.org/packages/fd/0d/5b87f5d7c3501bd5296b0bba7ba8a3eaf639ded53b9a17e910ee3363dfc0/usd_core-26.5-cp314-none-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:755c469ec762f3b69d87f5e8af8f8098e4c107bf4c15ce570a042ac2fc2dbb76", size = 28776483, upload-time = "2026-04-24T20:18:05.082Z" }, + { url = "https://files.pythonhosted.org/packages/5a/48/d29a4649df00455174a5979fc8291021199bb2115d623378364b58055bb5/usd_core-26.5-cp314-none-win_amd64.whl", hash = "sha256:7654b5dfef6e7177849aa7e69962feb82a5312ad08469983214aae5821601296", size = 14043860, upload-time = "2026-04-24T20:18:07.896Z" }, +] + [[package]] name = "uuid-utils" version = "0.14.0" From 643f7af3790d23620162348bda9e73465648f9c6 Mon Sep 17 00:00:00 2001 From: Nabla7 Date: Fri, 26 Jun 2026 04:43:26 +0200 Subject: [PATCH 03/73] Add G1 scene navigation visualization --- dimos/robot/unitree/g1/assets/g1_29dof.xml | 3 + .../blueprints/basic/unitree_g1_groot_wbc.py | 210 ++++++++++++---- dimos/robot/unitree/g1/g1_rerun.py | 18 ++ dimos/simulation/adapters/whole_body/g1.py | 2 +- dimos/simulation/engines/mujoco_engine.py | 225 ++++++++++++++++-- dimos/simulation/engines/mujoco_sim_module.py | 193 +++++++++++---- .../engines/test_mujoco_sim_module.py | 53 +++++ dimos/simulation/mujoco/depth_camera.py | 6 +- dimos/simulation/mujoco/entity_scene.py | 40 +++- dimos/simulation/scene_assets/spec.py | 12 - dimos/simulation/scenes/catalog.py | 31 +-- dimos/visualization/rerun/bridge.py | 30 ++- dimos/visualization/rerun/scene_package.py | 108 +++++++++ .../visualization/rerun/test_scene_package.py | 108 +++++++++ dimos/visualization/rerun/urdf_robot.py | 217 +++++++++++++++++ dimos/web/templates/rerun_dashboard.html | 2 +- 16 files changed, 1100 insertions(+), 158 deletions(-) create mode 100644 dimos/visualization/rerun/scene_package.py create mode 100644 dimos/visualization/rerun/test_scene_package.py create mode 100644 dimos/visualization/rerun/urdf_robot.py diff --git a/dimos/robot/unitree/g1/assets/g1_29dof.xml b/dimos/robot/unitree/g1/assets/g1_29dof.xml index b1dd123152..48edf7e822 100644 --- a/dimos/robot/unitree/g1/assets/g1_29dof.xml +++ b/dimos/robot/unitree/g1/assets/g1_29dof.xml @@ -154,6 +154,9 @@ + + + diff --git a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py index 1709930830..de3e2732d0 100644 --- a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py @@ -28,14 +28,16 @@ Sim (``--simulation``): MujocoSimModule (in-process MuJoCo + SHM) + sim_mujoco_g1 adapter. 50 Hz tick (matches the rate the policy was trained at). No arming - ramp, no dry-run, no servo_arms -- sim physics doesn't gravity-collapse - the arms between trajectories. + ramp and no dry-run. The 14 arm joints are still held with the same + lower-priority servo task as hardware so headless and viewer runs do not + depend on incidental startup timing. Usage: dimos run unitree-g1-groot-wbc # real hardware dimos --simulation mujoco run unitree-g1-groot-wbc # sim - dimos --simulation mujoco --scene data/scene_packages/dimos_office/scene.meta.json \\ - run unitree-g1-groot-wbc # sim in a cooked scene package + dimos --simulation mujoco --scene none run unitree-g1-groot-wbc + dimos --simulation mujoco --scene office run unitree-g1-groot-wbc + dimos --simulation mujoco --scene supermarket run unitree-g1-groot-wbc Overrides (replace the old env-var dance): dimos run unitree-g1-groot-wbc \\ @@ -61,11 +63,25 @@ from dimos.core.global_config import global_config from dimos.core.transport import LCMTransport from dimos.hardware.whole_body.spec import WholeBodyConfig +from dimos.mapping.costmapper import CostMapper +from dimos.mapping.pointclouds.occupancy import HeightCostConfig from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.nav_msgs.Path import Path as NavPath from dimos.msgs.sensor_msgs.Imu import Imu from dimos.msgs.sensor_msgs.JointState import JointState from dimos.msgs.sensor_msgs.MotorCommandArray import MotorCommandArray +from dimos.navigation.movement_manager.movement_manager import MovementManager +from dimos.navigation.replanning_a_star.module import ReplanningAStarPlanner +from dimos.robot.unitree.g1.config import G1 +from dimos.robot.unitree.g1.g1_rerun import ( + G1_RERUN_ROOT, + g1_costmap, + g1_urdf_joint_state, + g1_urdf_static_robot, +) from dimos.utils.data import LfsPath +from dimos.visualization.rerun.scene_package import scene_package_static_entities +from dimos.visualization.vis_module import vis_module # Lazy data handles. LfsPath only triggers the LFS pull on first # str()/open(); using ``get_data(...)`` at import time would block the @@ -77,11 +93,53 @@ _adapter_address: str | Path _cmd_vel_topic = "/cmd_vel" if global_config.simulation else "/g1/cmd_vel" +_MUJOCO_LIDAR_CAMERAS = ( + "lidar_front_camera", + "lidar_left_camera", + "lidar_right_camera", +) +_MUJOCO_LIDAR_CAMERA = _MUJOCO_LIDAR_CAMERAS[0] +# Robot geoms occupy groups 0/1. The legacy floor uses group 2, and cooked +# scene packages/entities use group 3, so lidar should render world geometry. +_MUJOCO_LIDAR_GEOM_GROUPS = (2, 3) +assert G1.height_clearance is not None and G1.width_clearance is not None +_MUJOCO_LIDAR_BASE_KWARGS: dict[str, Any] = { + "width": 320, + "height": 240, + "fps": 2, + "enable_color": False, + "enable_depth": False, + "enable_pointcloud": True, + "pointcloud_fps": 1.0, + "enable_mujoco_lidar": True, + "mujoco_lidar_geom_groups": list(_MUJOCO_LIDAR_GEOM_GROUPS), + "mujoco_lidar_raycast_width": 64, + "mujoco_lidar_raycast_height": 32, + "mujoco_lidar_robot_exclusion_radius": G1.width_clearance, +} +_G1_SUPERMARKET_COMPOSED_MJB = ( + "mujoco/composed/unitree-g1-groot-wbc_spawn_9p2_11p8_yaw_m1p57_static_only_lidar.mjb" +) +_G1_NAV_VOXEL_RESOLUTION = 0.05 +_G1_NAV_OVERHEAD_SAFETY_MARGIN = 0.2 +_G1_NAV_MAX_STEP_HEIGHT = 0.10 +_G1_NAV_ROTATION_DIAMETER = 0.8 +_G1_NAV_SAFE_RADIUS_MARGIN = 0.6 + + +def _mujoco_lidar_kwargs(camera_name: str, camera_names: tuple[str, ...]) -> dict[str, Any]: + return { + "camera_name": camera_name, + "mujoco_lidar_camera_names": list(camera_names), + **_MUJOCO_LIDAR_BASE_KWARGS, + } + if global_config.simulation and global_config.simulation != "mujoco": raise ValueError("unitree-g1-groot-wbc only supports --simulation mujoco") if global_config.simulation == "mujoco": + from dimos.mapping.voxels import VoxelGridMapper from dimos.simulation.engines.mujoco_sim_module import MujocoSimModule from dimos.simulation.engines.robot_sim_binding import ( RobotSimSpec, @@ -118,9 +176,7 @@ def _legacy_mujoco_backend() -> Any: address=_MJCF_PATH, headless=True, dof=29, - enable_color=False, - enable_depth=False, - enable_pointcloud=False, + **_mujoco_lidar_kwargs(_MUJOCO_LIDAR_CAMERA, _MUJOCO_LIDAR_CAMERAS), inject_legacy_assets=True, robot_sim_spec=_g1_sim_spec, ) @@ -129,6 +185,21 @@ def _scene_mujoco_backend() -> tuple[Any, str | Path]: if global_config.scene is None: return _legacy_mujoco_backend(), _MJCF_PATH + scene_path = Path(str(global_config.scene)).expanduser() + if scene_path.suffix.lower() == ".mjb": + if not scene_path.exists(): + raise FileNotFoundError(f"MuJoCo binary scene not found: {scene_path}") + return ( + MujocoSimModule.blueprint( + address=scene_path, + headless=True, + dof=29, + **_mujoco_lidar_kwargs(_MUJOCO_LIDAR_CAMERA, _MUJOCO_LIDAR_CAMERAS), + robot_sim_spec=_g1_sim_spec, + ), + scene_path, + ) + from dimos.simulation.scenes.catalog import resolve_scene_package package = resolve_scene_package(global_config.scene) @@ -137,6 +208,19 @@ def _scene_mujoco_backend() -> tuple[Any, str | Path]: if package.mujoco_scene_path is None: raise ValueError(f"scene package has no MuJoCo scene artifact: {package.metadata_path}") + composed_scene = _precomposed_g1_scene(package.package_dir) + if composed_scene is not None: + return ( + MujocoSimModule.blueprint( + address=composed_scene, + headless=True, + dof=29, + **_mujoco_lidar_kwargs(_MUJOCO_LIDAR_CAMERA, _MUJOCO_LIDAR_CAMERAS), + robot_sim_spec=_g1_sim_spec, + ), + composed_scene, + ) + return ( MujocoSimModule.blueprint( scene_xml=package.mujoco_scene_path, @@ -146,14 +230,18 @@ def _scene_mujoco_backend() -> tuple[Any, str | Path]: scene_entities=package.entities, headless=True, dof=29, - enable_color=False, - enable_depth=False, - enable_pointcloud=False, + **_mujoco_lidar_kwargs(_MUJOCO_LIDAR_CAMERA, _MUJOCO_LIDAR_CAMERAS), robot_sim_spec=_g1_sim_spec, ), _ROBOT_ONLY_MJCF_PATH, ) + def _precomposed_g1_scene(package_dir: Path) -> Path | None: + if "supermarket" not in package_dir.name: + return None + candidate = package_dir / _G1_SUPERMARKET_COMPOSED_MJB + return candidate if candidate.exists() else None + # Sim backend: MuJoCo engine via SHM. _backend, _adapter_address = _scene_mujoco_backend() # MujocoSimModule's ``odom`` Out is the sole producer of ``/odom`` @@ -166,9 +254,35 @@ def _scene_mujoco_backend() -> tuple[Any, str | Path]: _auto_dry_run = False _default_ramp_seconds = 0.0 _decimation: int | None = 1 - # Sim physics holds the arms between trajectories on its own -- no - # servo task needed. - _arm_holder: TaskConfig | None = None + _arm_holder = TaskConfig( + name="servo_arms", + type="servo", + joint_names=g1_arms, + priority=10, + auto_start=True, + params={"default_positions": ARM_DEFAULT_POSE}, + ) + _mapper = VoxelGridMapper.blueprint(emit_every=1) + _nav_stack = autoconnect( + _mapper, + CostMapper.blueprint( + config=HeightCostConfig( + resolution=_G1_NAV_VOXEL_RESOLUTION, + can_pass_under=G1.height_clearance + _G1_NAV_OVERHEAD_SAFETY_MARGIN, + can_climb=_G1_NAV_MAX_STEP_HEIGHT, + ), + initial_safe_radius_meters=G1.width_clearance + _G1_NAV_SAFE_RADIUS_MARGIN, + ), + ReplanningAStarPlanner.blueprint( + robot_width=G1.width_clearance, + robot_rotation_diameter=_G1_NAV_ROTATION_DIAMETER, + ), + MovementManager.blueprint(), + ) + _remappings = [ + (VoxelGridMapper, "lidar", "pointcloud"), + (ControlCoordinator, "twist_command", "cmd_vel"), + ] else: from dimos.robot.unitree.g1.wholebody_connection import G1WholeBodyConnection @@ -193,6 +307,8 @@ def _scene_mujoco_backend() -> tuple[Any, str | Path]: auto_start=True, params={"default_positions": ARM_DEFAULT_POSE}, ) + _nav_stack = MovementManager.blueprint() + _remappings = [(ControlCoordinator, "twist_command", "cmd_vel")] def _g1_groot_rerun_blueprint() -> Any: @@ -212,43 +328,46 @@ def _g1_groot_rerun_blueprint() -> Any: ) -def _static_g1_body(rr: Any) -> Any: - return rr.Boxes3D( - half_sizes=[0.25, 0.20, 0.6], - centers=[[0.0, 0.0, 0.6]], - colors=[(0, 255, 127)], - fill_mode="MajorWireframe", - ) +def _g1_nav_path(path: NavPath) -> Any: + return path.to_rerun(z_offset=0.3) + +_static_rerun_entities = { + # MujocoSimModule logs odom as a Transform3D at world/odom; the robot + # mesh lives underneath and link transforms are driven by joint state. + G1_RERUN_ROOT: g1_urdf_static_robot(root_path=G1_RERUN_ROOT), +} +_static_rerun_entities.update(scene_package_static_entities(global_config.scene)) _rerun_config = { "blueprint": _g1_groot_rerun_blueprint, - "static": { - # MujocoSimModule logs odom as a Transform3D at world/odom. - # This body marker inherits that transform, giving dimos-viewer - # a visible robot anchor until a richer joint/URDF view exists. - "world/odom/g1": _static_g1_body, + "visual_override": { + # This blueprint uses raycast lidar, so suppress raw camera streams + # in Rerun. + "world/color_image": None, + "world/camera_info": None, + "world/depth_image": None, + "world/depth_camera_info": None, + "world/coordinator_joint_state": g1_urdf_joint_state(root_path=G1_RERUN_ROOT), + "world/global_costmap": g1_costmap, + "world/navigation_costmap": g1_costmap, + "world/path": _g1_nav_path, }, + "max_hz": { + "world/coordinator_joint_state": 20.0, + "world/global_map": 1.0, + "world/global_costmap": 2.0, + "world/navigation_costmap": 2.0, + # The planner publishes an empty Path() immediately before the new + # planned path. Throttling this entity drops the real path. + "world/path": 0, + }, + "static": _static_rerun_entities, } def _viewer() -> Any: - if global_config.viewer == "none": - return autoconnect() - if global_config.viewer != "rerun": - raise ValueError(f"Unsupported viewer backend for G1 GR00T WBC: {global_config.viewer}") - - from dimos.visualization.rerun.bridge import RerunBridgeModule - from dimos.visualization.rerun.websocket_server import RerunWebSocketServer - - return autoconnect( - RerunBridgeModule.blueprint( - **_rerun_config, - rerun_open=global_config.rerun_open, - rerun_web=global_config.rerun_web, - ), - RerunWebSocketServer.blueprint(), - ) + return vis_module(viewer_backend=global_config.viewer, rerun_config=_rerun_config) _coordinator = ControlCoordinator.blueprint( @@ -284,8 +403,7 @@ def _viewer() -> Any: ).transports( { ("joint_command", JointState): LCMTransport("/g1/joint_command", JointState), - ("twist_command", Twist): LCMTransport(_cmd_vel_topic, Twist), - ("tele_cmd_vel", Twist): LCMTransport(_cmd_vel_topic, Twist), + ("cmd_vel", Twist): LCMTransport(_cmd_vel_topic, Twist), # Real-hw only: the transport_lcm adapter speaks to # G1WholeBodyConnection over these topics. autoconnect already # matches by (name, type) so sim doesn't need them -- they're @@ -296,6 +414,8 @@ def _viewer() -> Any: } ) -unitree_g1_groot_wbc = autoconnect(_backend, _coordinator, _viewer()).global_config( - robot_model="unitree_g1" +unitree_g1_groot_wbc = ( + autoconnect(_backend, _coordinator, _nav_stack, _viewer()) + .remappings(_remappings) + .global_config(robot_model="unitree_g1") ) diff --git a/dimos/robot/unitree/g1/g1_rerun.py b/dimos/robot/unitree/g1/g1_rerun.py index 8162813031..052e67063f 100644 --- a/dimos/robot/unitree/g1/g1_rerun.py +++ b/dimos/robot/unitree/g1/g1_rerun.py @@ -20,6 +20,14 @@ import numpy as np +from dimos.visualization.rerun.urdf_robot import ( + UrdfRobotJointStateRerunFactory, + UrdfRobotStaticRerunFactory, +) + +G1_RERUN_ROOT = "world/odom/g1" +G1_RERUN_URDF = "g1_urdf/g1.fixed.urdf" + # Classic costmap palette, indexed by grid value + 1: # transparent unknown, blue free, orange occupied, red lethal. _COSTMAP_LOOKUP_TABLE = np.zeros((102, 4), dtype=np.uint8) @@ -37,6 +45,16 @@ def g1_costmap(grid: Any) -> Any: return grid.to_rerun(color_lookup_table=_COSTMAP_LOOKUP_TABLE, z_offset=0.02) +def g1_urdf_static_robot(root_path: str = G1_RERUN_ROOT) -> UrdfRobotStaticRerunFactory: + """Create a static Rerun logger for the G1 URDF visual meshes.""" + return UrdfRobotStaticRerunFactory(urdf_path=G1_RERUN_URDF, root_path=root_path) + + +def g1_urdf_joint_state(root_path: str = G1_RERUN_ROOT) -> UrdfRobotJointStateRerunFactory: + """Create a Rerun JointState converter for the G1 URDF.""" + return UrdfRobotJointStateRerunFactory(urdf_path=G1_RERUN_URDF, root_path=root_path) + + def g1_static_robot(rr: Any) -> list[Any]: """Static G1 humanoid wireframe box attached to the sensor TF frame. diff --git a/dimos/simulation/adapters/whole_body/g1.py b/dimos/simulation/adapters/whole_body/g1.py index 43c8a98fb7..c925db27d0 100644 --- a/dimos/simulation/adapters/whole_body/g1.py +++ b/dimos/simulation/adapters/whole_body/g1.py @@ -48,7 +48,7 @@ _NUM_MOTORS = 29 -_READY_WAIT_TIMEOUT_S = 60.0 +_READY_WAIT_TIMEOUT_S = 180.0 _READY_WAIT_POLL_S = 0.1 _ATTACH_RETRY_TIMEOUT_S = 30.0 _ATTACH_RETRY_POLL_S = 0.2 diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index 0454a6bda1..a07fca1285 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -51,12 +51,43 @@ _RESET_WAIT_TIMEOUT_S = 5.0 +def _camera_name_candidates(camera_name: str) -> tuple[str, ...]: + stripped = camera_name.strip("/") + candidates = [camera_name] + if stripped != camera_name: + candidates.append(stripped) + elif stripped: + candidates.append(f"/{stripped}") + return tuple(dict.fromkeys(candidates)) + + +def _camera_ray_directions(width: int, height: int, fovy_degrees: float) -> NDArray[np.float64]: + """Return normalized camera-frame ray directions for MuJoCo cameras. + + MuJoCo/OpenGL cameras look along local -Z, with +X right and +Y up. + """ + fovy = math.radians(fovy_degrees) + focal = height / (2.0 * math.tan(fovy / 2.0)) + cx = width / 2.0 + cy = height / 2.0 + + ys, xs = np.mgrid[0:height, 0:width] + x = (xs + 0.5 - cx) / focal + y = -(ys + 0.5 - cy) / focal + z = -np.ones_like(x) + directions = np.stack((x, y, z), axis=-1).reshape(-1, 3).astype(np.float64) + norms = np.linalg.norm(directions, axis=1, keepdims=True) + return directions / norms + + @dataclass class CameraConfig: name: str width: int = 640 height: int = 480 fps: float = 15.0 + max_geom: int | None = 10000 + geom_groups: tuple[int, ...] | None = None @dataclass @@ -69,16 +100,46 @@ class CameraFrame: timestamp: float +@dataclass +class RaycastLidarConfig: + name: str + width: int = 64 + height: int = 32 + fps: float = 1.0 + min_range: float = 0.2 + max_range: float = 3.0 + max_height: float = 1.2 + geom_groups: tuple[int, ...] | None = None + robot_exclusion_radius: float = 0.0 + + +@dataclass +class RaycastLidarFrame: + points: NDArray[np.float32] + timestamp: float + + @dataclass class _CameraRendererState: cfg: CameraConfig cam_id: int rgb_renderer: mujoco.Renderer depth_renderer: mujoco.Renderer + scene_option: mujoco.MjvOption | None interval: float last_render_time: float = 0.0 +@dataclass +class _RaycastLidarState: + cfg: RaycastLidarConfig + cam_id: int + ray_directions_camera: NDArray[np.float64] + geomgroup: NDArray[np.uint8] + interval: float + last_cast_time: float = 0.0 + + class MujocoEngine(SimulationEngine): """ MuJoCo simulation engine. @@ -93,6 +154,7 @@ def __init__( config_path: Path, headless: bool, cameras: list[CameraConfig] | None = None, + raycast_lidars: list[RaycastLidarConfig] | None = None, on_before_step: StepHook | None = None, on_after_step: StepHook | None = None, assets: dict[str, bytes] | None = None, @@ -111,7 +173,8 @@ def __init__( self._spawn_yaw = spawn_yaw self._reset_joint_positions = reset_joint_positions - xml_path = self._resolve_xml_path(config_path) + model_path = self._resolve_model_path(config_path) + binary_model = model_path.suffix.lower() == ".mjb" if model is not None: if assets is not None: raise ValueError("MujocoEngine cannot use injected assets with a precompiled model") @@ -120,15 +183,19 @@ def __init__( # MJCFs that reference meshes by bare filename (e.g. menagerie # G1) need the mesh bytes injected by name; from_xml_path can't # find them on disk. - with open(xml_path) as f: + if binary_model: + raise ValueError("MujocoEngine cannot inject XML assets into a binary MJB model") + with open(model_path) as f: xml_str = f.read() self._model = mujoco.MjModel.from_xml_string(xml_str, assets=assets) + elif binary_model: + self._model = mujoco.MjModel.from_binary_path(str(model_path)) else: - self._model = mujoco.MjModel.from_xml_path(str(xml_path)) - self._xml_path = xml_path + self._model = mujoco.MjModel.from_xml_path(str(model_path)) + self._xml_path = model_path self._data = mujoco.MjData(self._model) - mapping_xml_path = None if model is not None else self._xml_path + mapping_xml_path = None if model is not None or binary_model else self._xml_path self._joint_mappings = build_joint_mappings(mapping_xml_path, self._model) self._robot_binding: RobotSimBinding | None = None if robot_sim_spec is not None: @@ -171,6 +238,9 @@ def __init__( self._camera_configs = cameras or [] self._camera_frames: dict[str, CameraFrame] = {} self._camera_lock = threading.Lock() + self._raycast_lidar_configs = raycast_lidars or [] + self._raycast_lidar_frames: dict[str, RaycastLidarFrame] = {} + self._raycast_lidar_lock = threading.Lock() def set_step_hooks( self, @@ -185,14 +255,14 @@ def set_step_hooks( self._on_before_step = before self._on_after_step = after - def _resolve_xml_path(self, config_path: Path) -> Path: + def _resolve_model_path(self, config_path: Path) -> Path: if config_path is None: raise ValueError("config_path is required for MuJoCo simulation loading") resolved = config_path.expanduser() - xml_path = resolved / "scene.xml" if resolved.is_dir() else resolved - if not xml_path.exists(): - raise FileNotFoundError(f"MuJoCo XML not found: {xml_path}") - return xml_path + model_path = resolved / "scene.xml" if resolved.is_dir() else resolved + if not model_path.exists(): + raise FileNotFoundError(f"MuJoCo model not found: {model_path}") + return model_path def _find_first_freejoint_adrs(self) -> tuple[int | None, int | None]: if self._model.njnt > 0 and int(self._model.jnt_type[0]) == _MJJNT_FREE: @@ -291,23 +361,76 @@ def _init_cameras(self) -> dict[str, _CameraRendererState]: """Create renderers for all configured cameras""" cam_renderers: dict[str, _CameraRendererState] = {} for cfg in self._camera_configs: - cam_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_CAMERA, cfg.name) + cam_id = self._camera_id(cfg.name) if cam_id < 0: logger.warning("Camera not found in MJCF, skipping", camera_name=cfg.name) continue - rgb_renderer = mujoco.Renderer(self._model, height=cfg.height, width=cfg.width) - depth_renderer = mujoco.Renderer(self._model, height=cfg.height, width=cfg.width) + max_geom = cfg.max_geom + if max_geom is None: + max_geom = max(10000, int(self._model.ngeom) + 1000) + rgb_renderer = mujoco.Renderer( + self._model, + height=cfg.height, + width=cfg.width, + max_geom=max_geom, + ) + depth_renderer = mujoco.Renderer( + self._model, + height=cfg.height, + width=cfg.width, + max_geom=max_geom, + ) depth_renderer.enable_depth_rendering() + scene_option = None + if cfg.geom_groups is not None: + scene_option = mujoco.MjvOption() + scene_option.geomgroup[:] = 0 + for group in cfg.geom_groups: + if 0 <= group < len(scene_option.geomgroup): + scene_option.geomgroup[group] = 1 interval = 1.0 / cfg.fps if cfg.fps > 0 else float("inf") cam_renderers[cfg.name] = _CameraRendererState( cfg=cfg, cam_id=cam_id, rgb_renderer=rgb_renderer, depth_renderer=depth_renderer, + scene_option=scene_option, interval=interval, ) return cam_renderers + def _init_raycast_lidars(self) -> dict[str, _RaycastLidarState]: + lidar_states: dict[str, _RaycastLidarState] = {} + for cfg in self._raycast_lidar_configs: + cam_id = self._camera_id(cfg.name) + if cam_id < 0: + logger.warning( + "Raycast lidar camera not found in MJCF, skipping", + camera_name=cfg.name, + ) + continue + + geomgroup = np.zeros(6, dtype=np.uint8) + if cfg.geom_groups is None: + geomgroup[:] = 1 + else: + for group in cfg.geom_groups: + if 0 <= group < len(geomgroup): + geomgroup[group] = 1 + + lidar_states[cfg.name] = _RaycastLidarState( + cfg=cfg, + cam_id=cam_id, + ray_directions_camera=_camera_ray_directions( + cfg.width, + cfg.height, + float(self._model.cam_fovy[cam_id]), + ), + geomgroup=geomgroup, + interval=1.0 / cfg.fps if cfg.fps > 0 else float("inf"), + ) + return lidar_states + def _render_cameras(self, now: float, cam_renderers: dict[str, _CameraRendererState]) -> None: """Render all due cameras and store frames. Must be called from sim thread.""" for state in cam_renderers.values(): @@ -315,10 +438,14 @@ def _render_cameras(self, now: float, cam_renderers: dict[str, _CameraRendererSt continue state.last_render_time = now - state.rgb_renderer.update_scene(self._data, camera=state.cam_id) + state.rgb_renderer.update_scene( + self._data, camera=state.cam_id, scene_option=state.scene_option + ) rgb = state.rgb_renderer.render().copy() - state.depth_renderer.update_scene(self._data, camera=state.cam_id) + state.depth_renderer.update_scene( + self._data, camera=state.cam_id, scene_option=state.scene_option + ) depth = state.depth_renderer.render().copy() frame = CameraFrame( @@ -332,6 +459,58 @@ def _render_cameras(self, now: float, cam_renderers: dict[str, _CameraRendererSt with self._camera_lock: self._camera_frames[state.cfg.name] = frame + def _raycast_lidars( + self, + now: float, + lidar_states: dict[str, _RaycastLidarState], + ) -> None: + """Raycast lidar frames from MuJoCo cameras. Must be called from sim thread.""" + bodyexclude = self._robot_binding.root_body_id if self._robot_binding is not None else -1 + for state in lidar_states.values(): + if now - state.last_cast_time < state.interval: + continue + state.last_cast_time = now + + origin = self._data.cam_xpos[state.cam_id].copy() + camera_mat = self._data.cam_xmat[state.cam_id].reshape(3, 3).copy() + directions_world = state.ray_directions_camera @ camera_mat.T + n_rays = directions_world.shape[0] + geom_ids = np.full(n_rays, -1, dtype=np.int32) + distances = np.full(n_rays, -1.0, dtype=np.float64) + mujoco.mj_multiRay( + self._model, + self._data, + origin, + directions_world.ravel(), + state.geomgroup, + 1, + bodyexclude, + geom_ids, + distances, + None, + n_rays, + state.cfg.max_range, + ) + valid = (distances >= state.cfg.min_range) & (distances <= state.cfg.max_range) + valid &= np.abs(state.ray_directions_camera[:, 1] * distances) <= state.cfg.max_height + if np.any(valid): + points_world = origin + directions_world[valid] * distances[valid, None] + if state.cfg.robot_exclusion_radius > 0.0 and self._root_qpos_adr is not None: + root_xy = self._data.qpos[self._root_qpos_adr : self._root_qpos_adr + 2] + keep = ( + np.linalg.norm(points_world[:, :2] - root_xy, axis=1) + >= state.cfg.robot_exclusion_radius + ) + points_world = points_world[keep] + arr = points_world.astype(np.float32) + else: + arr = np.empty((0, 3), dtype=np.float32) + with self._raycast_lidar_lock: + self._raycast_lidar_frames[state.cfg.name] = RaycastLidarFrame( + points=arr, + timestamp=now, + ) + @staticmethod def _close_cam_renderers(cam_renderers: dict[str, _CameraRendererState]) -> None: for state in cam_renderers.values(): @@ -401,6 +580,7 @@ def _sim_loop(self) -> None: # Camera renderers: created once in the sim thread cam_renderers = self._init_cameras() + lidar_states = self._init_raycast_lidars() def _step_once(sync_viewer: bool) -> None: loop_start = time.time() @@ -429,6 +609,7 @@ def _step_once(sync_viewer: bool) -> None: except Exception as exc: logger.error("on_after_step failed", error=str(exc)) self._render_cameras(loop_start, cam_renderers) + self._raycast_lidars(loop_start, lidar_states) elapsed = time.time() - loop_start sleep_time = dt - elapsed @@ -689,9 +870,21 @@ def read_camera(self, camera_name: str) -> CameraFrame | None: with self._camera_lock: return self._camera_frames.get(camera_name) + def read_raycast_lidar(self, camera_name: str) -> RaycastLidarFrame | None: + """Read the latest raycast lidar frame for a camera (thread-safe).""" + with self._raycast_lidar_lock: + return self._raycast_lidar_frames.get(camera_name) + def get_camera_fovy(self, camera_name: str) -> float | None: """Get vertical field of view for a named camera, in degrees.""" - cam_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_CAMERA, camera_name) + cam_id = self._camera_id(camera_name) if cam_id < 0: return None return float(self._model.cam_fovy[cam_id]) + + def _camera_id(self, camera_name: str) -> int: + for candidate in _camera_name_candidates(camera_name): + cam_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_CAMERA, candidate) + if cam_id >= 0: + return int(cam_id) + return -1 diff --git a/dimos/simulation/engines/mujoco_sim_module.py b/dimos/simulation/engines/mujoco_sim_module.py index 82c4a60876..480a2f3361 100644 --- a/dimos/simulation/engines/mujoco_sim_module.py +++ b/dimos/simulation/engines/mujoco_sim_module.py @@ -56,6 +56,7 @@ CameraConfig, CameraFrame, MujocoEngine, + RaycastLidarConfig, ) from dimos.simulation.engines.mujoco_shm import ( CMD_MODE_PD_TAU, @@ -63,6 +64,7 @@ shm_key_from_path, ) from dimos.simulation.engines.robot_sim_binding import RobotSimSpec +from dimos.simulation.mujoco.constants import LIDAR_RESOLUTION, MAX_HEIGHT, MAX_RANGE, MIN_RANGE from dimos.spec import perception from dimos.utils.logging_config import setup_logger @@ -239,6 +241,18 @@ class MujocoSimModuleConfig(ModuleConfig, DepthCameraConfig): enable_pointcloud: bool = False pointcloud_fps: float = 5.0 camera_info_fps: float = 1.0 + # Optional MuJoCo-native lidar: cast rays from one or more named cameras + # and publish world-frame PointCloud2 points on ``pointcloud``. + enable_mujoco_lidar: bool = False + mujoco_lidar_camera_names: list[str] = Field(default_factory=list) + mujoco_lidar_voxel_size: float = LIDAR_RESOLUTION + mujoco_lidar_geom_groups: list[int] = Field(default_factory=lambda: [0, 1, 2, 3, 4, 5]) + mujoco_lidar_raycast_width: int = 64 + mujoco_lidar_raycast_height: int = 32 + mujoco_lidar_min_range: float = MIN_RANGE + mujoco_lidar_max_range: float = MAX_RANGE + mujoco_lidar_max_height: float = MAX_HEIGHT + mujoco_lidar_robot_exclusion_radius: float = 0.0 # Inject menagerie/dimos-bundled mesh bytes (via # dimos.simulation.mujoco.model.get_assets) into MjModel.from_xml_string. # MJCFs that reference meshes by bare filename (G1 GR00T, Go2) need this; @@ -375,32 +389,68 @@ def start(self) -> None: from dimos.simulation.mujoco.model import get_assets engine_assets = get_assets() - # Compose the camera list. Each registered camera blocks the - # sim thread inside _step_once (mujoco_engine._render_cameras - # does update_scene + GPU render synchronously between physics - # steps - typically 5-30 ms per camera), so registering a camera - # nobody consumes burns the 500 Hz tick deadline for nothing. - # Skip the primary camera entirely when none of color / depth / - # pointcloud is enabled. - cameras: list[CameraConfig] = [] + # Compose rendered cameras separately from raycast lidar. Each + # rendered camera blocks the sim thread, so MuJoCo lidar does not + # register depth cameras here. + cameras_by_name: dict[str, CameraConfig] = {} + raycast_lidars: list[RaycastLidarConfig] = [] + + def add_camera( + name: str, + *, + geom_groups: list[int] | None = None, + max_geom: int | None = 10000, + ) -> None: + if not name: + return + groups = tuple(geom_groups) if geom_groups is not None else None + existing = cameras_by_name.get(name) + if existing is not None: + if groups is not None: + existing.geom_groups = groups + existing.max_geom = max_geom + return + cameras_by_name[name] = CameraConfig( + name=name, + width=self.config.width, + height=self.config.height, + fps=float(self.config.fps), + max_geom=max_geom, + geom_groups=groups, + ) + primary_needed = ( - self.config.enable_color or self.config.enable_depth or self.config.enable_pointcloud + self.config.enable_color + or self.config.enable_depth + or (self.config.enable_pointcloud and not self.config.enable_mujoco_lidar) ) if primary_needed: - cameras.append( - CameraConfig( - name=self.config.camera_name, - width=self.config.width, - height=self.config.height, - fps=float(self.config.fps), + add_camera(self.config.camera_name) + + if self.config.enable_pointcloud and self.config.enable_mujoco_lidar: + for camera_name in self._mujoco_lidar_camera_names(): + raycast_lidars.append( + RaycastLidarConfig( + name=camera_name, + width=self.config.mujoco_lidar_raycast_width, + height=self.config.mujoco_lidar_raycast_height, + fps=self.config.pointcloud_fps, + min_range=self.config.mujoco_lidar_min_range, + max_range=self.config.mujoco_lidar_max_range, + max_height=self.config.mujoco_lidar_max_height, + geom_groups=tuple(self.config.mujoco_lidar_geom_groups), + robot_exclusion_radius=self.config.mujoco_lidar_robot_exclusion_radius, + ) ) - ) + + cameras = list(cameras_by_name.values()) # Hooks are installed via set_step_hooks() after gripper detection # below, since they depend on the resolved gripper index. engine_kwargs: dict[str, Any] = dict( headless=self.config.headless, cameras=cameras, + raycast_lidars=raycast_lidars, robot_sim_spec=self.config.robot_sim_spec, reset_joint_positions=self.config.reset_joint_positions, ) @@ -495,8 +545,12 @@ def start(self) -> None: ) ) - # Optional pointcloud generation: back-projects primary camera depth. - if self.config.enable_pointcloud and self.config.enable_depth: + # Optional pointcloud generation. Default mode preserves the legacy + # RGB-D camera pointcloud; MuJoCo lidar mode publishes native raycast + # hits in world coordinates. + if self.config.enable_pointcloud and ( + self.config.enable_depth or self.config.enable_mujoco_lidar + ): pc_interval = 1.0 / self.config.pointcloud_fps self.register_disposable( rx.interval(pc_interval).subscribe( @@ -517,36 +571,57 @@ def start(self) -> None: def _compose_model(self) -> mujoco.MjModel: """Compose optional scene package MJCF + robot MJCF + entities.""" - from dimos.simulation.mujoco.entity_scene import add_entities_to_spec + from dimos.simulation.mujoco.entity_scene import add_entities_to_spec, spawn_penetrators if self.config.robot_mjcf is None: raise RuntimeError("MujocoSimModule: robot_mjcf is required for composition") - if self.config.scene_xml is not None: - spec_scene = mujoco.MjSpec.from_file(str(self.config.scene_xml)) - else: - spec_scene = mujoco.MjSpec() - - spec_robot = mujoco.MjSpec.from_file(str(self.config.robot_mjcf)) - if self.config.robot_meshdir is not None: - spec_robot.meshdir = str(self.config.robot_meshdir) - - # Keep the robot controller timing stable when attached to a scene - # package whose wrapper may have different default options. - spec_scene.option.timestep = spec_robot.option.timestep + def build_spec(*, force_static: frozenset[str] = frozenset()) -> mujoco.MjSpec: + if self.config.scene_xml is not None: + spec_scene = mujoco.MjSpec.from_file(str(self.config.scene_xml)) + else: + spec_scene = mujoco.MjSpec() + + spec_robot = mujoco.MjSpec.from_file(str(self.config.robot_mjcf)) + if self.config.robot_meshdir is not None: + spec_robot.meshdir = str(self.config.robot_meshdir) + + # Keep the robot controller timing stable when attached to a scene + # package whose wrapper may have different default options. + spec_scene.option.timestep = spec_robot.option.timestep + + spawn_xy = self.config.spawn_xy or (0.0, 0.0) + spawn_z = self.config.spawn_z if self.config.spawn_z is not None else 0.0 + frame_kwargs: dict[str, Any] = { + "pos": [float(spawn_xy[0]), float(spawn_xy[1]), float(spawn_z)], + } + if self.config.spawn_yaw is not None: + yaw = float(self.config.spawn_yaw) + frame_kwargs["quat"] = [math.cos(yaw * 0.5), 0.0, 0.0, math.sin(yaw * 0.5)] + frame = spec_scene.worldbody.add_frame( + **frame_kwargs, + ) + prefix = f"{self.config.robot_id}-" if self.config.robot_id else None + spec_scene.attach(spec_robot, prefix=prefix, frame=frame) - spawn_xy = self.config.spawn_xy or (0.0, 0.0) - spawn_z = self.config.spawn_z if self.config.spawn_z is not None else 0.0 - frame = spec_scene.worldbody.add_frame( - pos=[float(spawn_xy[0]), float(spawn_xy[1]), float(spawn_z)], - ) - prefix = f"{self.config.robot_id}-" if self.config.robot_id else None - spec_scene.attach(spec_robot, prefix=prefix, frame=frame) + if self.config.scene_entities: + add_entities_to_spec( + spec_scene, self.config.scene_entities, force_static=force_static + ) + return spec_scene + spec_scene = build_spec() + model = spec_scene.compile() if self.config.scene_entities: - add_entities_to_spec(spec_scene, self.config.scene_entities) - - return spec_scene.compile() + penetrators = spawn_penetrators(model) + if penetrators: + logger.warning( + "MujocoSimModule: scene entities spawn in deep contact; welding static", + count=len(penetrators), + samples=sorted(penetrators)[:20], + ) + model = build_spec(force_static=penetrators).compile() + return model @rpc def stop(self) -> None: @@ -877,6 +952,9 @@ def _publish_tf(self, ts: float, frame: CameraFrame | None) -> None: def _generate_pointcloud(self) -> None: if self._engine is None: return + if self.config.enable_mujoco_lidar: + self._generate_mujoco_lidar_pointcloud() + return # Back-project the primary camera's depth image. if self._camera_info_base is None: return @@ -906,3 +984,36 @@ def _generate_pointcloud(self) -> None: self.pointcloud.publish(pcd) except Exception as exc: logger.error("Pointcloud generation error", error=str(exc)) + + def _mujoco_lidar_camera_names(self) -> list[str]: + names = self.config.mujoco_lidar_camera_names + return list(names) if names else [self.config.camera_name] + + def _generate_mujoco_lidar_pointcloud(self) -> None: + engine = self._engine + if engine is None: + return + + all_points: list[NDArray[Any]] = [] + latest_ts = 0.0 + for camera_name in self._mujoco_lidar_camera_names(): + frame = engine.read_raycast_lidar(camera_name) + if frame is None: + continue + if frame.points.size > 0: + all_points.append(frame.points) + latest_ts = max(latest_ts, frame.timestamp) + + if not all_points: + return + + try: + pcd = PointCloud2.from_numpy( + np.vstack(all_points), + frame_id="world", + timestamp=latest_ts or time.time(), + ) + pcd = pcd.voxel_downsample(self.config.mujoco_lidar_voxel_size) + self.pointcloud.publish(pcd) + except Exception as exc: + logger.error("MuJoCo lidar pointcloud generation error", error=str(exc)) diff --git a/dimos/simulation/engines/test_mujoco_sim_module.py b/dimos/simulation/engines/test_mujoco_sim_module.py index 0074493858..f0fb1eb25c 100644 --- a/dimos/simulation/engines/test_mujoco_sim_module.py +++ b/dimos/simulation/engines/test_mujoco_sim_module.py @@ -216,6 +216,31 @@ def _scene_entity(entity_id: str) -> dict[str, object]: } +def _write_hull_obj(path: Path) -> None: + path.write_text( + """ +v 0 0 0 +v 0.1 0 0 +v 0 0.1 0 +v 0 0 0.1 +f 1 2 3 +f 1 2 4 +f 1 3 4 +f 2 3 4 +""".strip() + ) + + +def _mesh_scene_entity(entity_id: str, hull_path: Path) -> dict[str, object]: + entity = _scene_entity(entity_id) + descriptor = dict(entity["descriptor"]) # type: ignore[arg-type] + descriptor["shape_hint"] = "mesh" + descriptor["extents"] = [] + entity["descriptor"] = descriptor + entity["collision_paths"] = [str(hull_path)] + return entity + + @pytest.mark.mujoco def test_compose_model_attaches_robot_before_scene_entities(tmp_path: Path) -> None: mujoco = pytest.importorskip("mujoco") @@ -257,6 +282,34 @@ def test_compose_model_attaches_robot_before_scene_entities(tmp_path: Path) -> N assert any(name.endswith("hinge") for name in engine.joint_names) +@pytest.mark.mujoco +def test_compose_model_reuses_entity_mesh_assets(tmp_path: Path) -> None: + mujoco = pytest.importorskip("mujoco") + scene_xml = tmp_path / "scene.xml" + robot_xml = tmp_path / "robot.xml" + hull_obj = tmp_path / "shared_hull.obj" + _write_scene_xml(scene_xml) + _write_robot_xml(robot_xml) + _write_hull_obj(hull_obj) + + module = object.__new__(MujocoSimModule) + module.config = MujocoSimModuleConfig( + scene_xml=scene_xml, + robot_mjcf=robot_xml, + scene_entities=[ + _mesh_scene_entity("box_000", hull_obj), + _mesh_scene_entity("box_001", hull_obj), + ], + spawn_xy=(0.0, 0.0), + ) + + model = MujocoSimModule._compose_model(module) + + assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "entity:box_000") >= 0 + assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "entity:box_001") >= 0 + assert model.nmesh == 1 + + @pytest.mark.mujoco def test_engine_request_reset_to_applies_pose_in_sim_loop(tmp_path: Path) -> None: robot_xml = tmp_path / "freejoint.xml" diff --git a/dimos/simulation/mujoco/depth_camera.py b/dimos/simulation/mujoco/depth_camera.py index 486b740ffd..c7d9b4dd3c 100644 --- a/dimos/simulation/mujoco/depth_camera.py +++ b/dimos/simulation/mujoco/depth_camera.py @@ -58,7 +58,11 @@ def depth_image_to_point_cloud( o3d_depth = o3d.geometry.Image(depth_image.astype(np.float32)) # Create point cloud from depth image using Open3D - o3d_cloud = o3d.geometry.PointCloud.create_from_depth_image(o3d_depth, cam_intrinsics) + o3d_cloud = o3d.geometry.PointCloud.create_from_depth_image( + o3d_depth, + cam_intrinsics, + depth_scale=1.0, + ) # Convert Open3D point cloud to numpy array camera_points: NDArray[Any] = np.asarray(o3d_cloud.points) diff --git a/dimos/simulation/mujoco/entity_scene.py b/dimos/simulation/mujoco/entity_scene.py index c28c82c053..fa971d9d1f 100644 --- a/dimos/simulation/mujoco/entity_scene.py +++ b/dimos/simulation/mujoco/entity_scene.py @@ -75,6 +75,17 @@ def entity_body_name(entity_id: str) -> str: return f"{ENTITY_BODY_PREFIX}{entity_id}" +def _entity_id_from_geom_name(name: str) -> str | None: + if not name.startswith(ENTITY_BODY_PREFIX): + return None + rest = name[len(ENTITY_BODY_PREFIX) :] + marker = ":geom" + marker_idx = rest.rfind(marker) + if marker_idx < 0: + return None + return rest[:marker_idx] + + def _initial_entities(entities: list[dict[str, Any]]) -> list[dict[str, Any]]: return [e for e in entities if e.get("spawn", "initial") == "initial"] @@ -139,6 +150,12 @@ def _entity_collision_hulls(entity: dict[str, Any]) -> list[Path]: return paths +def _entity_mesh_asset_name(path: Path) -> str: + safe_stem = "".join(c if c.isalnum() else "_" for c in path.stem).strip("_") + digest = hashlib.sha256(str(path).encode()).hexdigest()[:12] + return f"entity_hull_{safe_stem}_{digest}" + + def _entity_friction(entity: dict[str, Any]) -> tuple[float, float, float]: """``physics.friction`` from entity metadata (scalar sliding or full [sliding, torsional, rolling] triple), else the scoot-able default.""" @@ -177,6 +194,7 @@ def add_entities_to_spec( """ import mujoco + mesh_assets: dict[Path, str] = {} for entity in _initial_entities(entities): descriptor = entity.get("descriptor", {}) entity_id = descriptor.get("entity_id") or entity.get("id") @@ -240,8 +258,12 @@ def add_entities_to_spec( body.mass = mass geom_kwargs.pop("mass", None) for i, hull_obj in enumerate(hull_paths): - mesh_name = f"{base_name}:hull{i:03d}" - spec.add_mesh(name=mesh_name, file=str(hull_obj)) + hull_path = hull_obj.resolve() + mesh_name = mesh_assets.get(hull_path) + if mesh_name is None: + mesh_name = _entity_mesh_asset_name(hull_path) + spec.add_mesh(name=mesh_name, file=str(hull_path)) + mesh_assets[hull_path] = mesh_name # Name geoms uniquely so MuJoCo's compile doesn't reject # duplicate-name collisions across multi-hull entities. gk = dict(geom_kwargs) @@ -289,8 +311,9 @@ def spawn_penetrators(model: mujoco.MjModel) -> frozenset[str]: continue for geom_id in (contact.geom1, contact.geom2): name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, int(geom_id)) or "" - if name.startswith(ENTITY_BODY_PREFIX) and name.endswith(":geom"): - bad.add(name[len(ENTITY_BODY_PREFIX) : -len(":geom")]) + entity_id = _entity_id_from_geom_name(name) + if entity_id is not None: + bad.add(entity_id) return frozenset(bad) @@ -336,12 +359,3 @@ def compose_entity_model(scene_package: ScenePackage) -> Path | None: mjb_path = out_dir / f"entities_{key}.mjb" mujoco.mj_saveModel(model, str(mjb_path)) return mjb_path - - -__all__ = [ - "ENTITY_BODY_PREFIX", - "add_entities_to_spec", - "compose_entity_model", - "entity_body_name", - "spawn_penetrators", -] diff --git a/dimos/simulation/scene_assets/spec.py b/dimos/simulation/scene_assets/spec.py index e24882ee48..35b666541d 100644 --- a/dimos/simulation/scene_assets/spec.py +++ b/dimos/simulation/scene_assets/spec.py @@ -262,15 +262,3 @@ def rewrite(value: Any) -> Any: if isinstance(artifacts, dict): for key, value in list(artifacts.items()): artifacts[key] = rewrite(value) - - -__all__ = [ - "ARTIFACT_FRAMES", - "BrowserCollisionSpec", - "BrowserVisualSpec", - "MujocoSceneSpec", - "SceneCookSpec", - "SceneMeshAlignment", - "ScenePackage", - "load_scene_package", -] diff --git a/dimos/simulation/scenes/catalog.py b/dimos/simulation/scenes/catalog.py index 13f56100e8..3a751e52f6 100644 --- a/dimos/simulation/scenes/catalog.py +++ b/dimos/simulation/scenes/catalog.py @@ -22,37 +22,17 @@ from dimos.simulation.scene_assets.spec import ScenePackage, load_scene_package from dimos.utils.data import get_data -DEFAULT_SCENE = "dimos-office" +DEFAULT_SCENE = "office" _DISABLED_SCENE_NAMES = {"", "none", "off", "disabled", "false", "0"} _ALIASES = { - "default": DEFAULT_SCENE, "office": DEFAULT_SCENE, "dimos-office": DEFAULT_SCENE, "dimos_office": DEFAULT_SCENE, - "office-splat": "dimos-office-splat", - "dimos-office-splat": "dimos-office-splat", - "dimos_office_splat": "dimos-office-splat", - "street": "street-lite", - "street-lite": "street-lite", - "street_lite": "street-lite", - "mall": "mall-babylon-nolights", - "mall_babylon_nolights": "mall-babylon-nolights", - "mall-babylon-nolights": "mall-babylon-nolights", - "lowpoly": "lowpoly-tdm", - "lowpoly-tdm": "lowpoly-tdm", - "lowpoly_tdm": "lowpoly-tdm", - "tdm": "lowpoly-tdm", - "apartment": "dimos-apartment", - "dimos-apartment": "dimos-apartment", - "dimos_apartment": "dimos-apartment", + "supermarket": "supermarket", } _PACKAGE_DIRS = { DEFAULT_SCENE: "dimos_office", - "dimos-office-splat": "dimos_office_splat", - "street-lite": "street_lite", - "mall-babylon-nolights": "mall_babylon_nolights", - "lowpoly-tdm": "lowpoly_tdm", - "dimos-apartment": "dimos_apartment", + "supermarket": "supermarket", } @@ -67,7 +47,7 @@ def resolve_scene_package( runtime composer attaches the robot via ``MjSpec.attach()``. """ if scene is None: - scene = DEFAULT_SCENE + return None scene_text = str(scene).strip() if scene_text.lower() in _DISABLED_SCENE_NAMES: @@ -129,6 +109,3 @@ def _resolve_dimos_office() -> ScenePackage: def _scene_package_dir() -> Path: return get_data("scene_packages") - - -__all__ = ["DEFAULT_SCENE", "resolve_scene_package"] diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 2f5fb1efa9..5c2aaa2e52 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -134,6 +134,12 @@ def _hex_to_rgba(hex_color: str) -> int: return int(h[:8], 16) +def _set_rerun_message_time(msg: Any) -> None: + ts = getattr(msg, "ts", None) + if isinstance(ts, (int, float)) and ts > 0: + rr.set_time("dimos_time", timestamp=float(ts)) + + def _with_graph_tab(bp: Blueprint) -> Blueprint: """Add a Graph tab alongside the existing viewer layout without changing it.""" @@ -169,7 +175,7 @@ class Config(ModuleConfig): visual_override: dict[Glob | str, Callable[[Any], Archetype] | None] = field( default_factory=dict ) - static: dict[str, Callable[[Any], Archetype]] = field(default_factory=dict) + static: dict[str, Callable[[Any], Any]] = field(default_factory=dict) max_hz: dict[str, float] = field(default_factory=dict) entity_prefix: str = "world" @@ -289,6 +295,8 @@ def _on_message(self, msg: Any, topic: Any) -> None: if not rerun_data: return + _set_rerun_message_time(msg) + # TFMessage for example returns list of (entity_path, archetype) tuples if is_rerun_multi(rerun_data): for path, archetype in rerun_data: @@ -433,10 +441,30 @@ def _log_connect_hints(self, grpc_port: int) -> None: def _log_static(self) -> None: for entity_path, factory in self.config.static.items(): data = factory(rr) + if is_rerun_multi(data): + logger.info( + "Rerun static entity", + entity_path=entity_path, + archetypes=[type(archetype).__name__ for _, archetype in data], + ) + for path, archetype in data: + rr.log(path, archetype, static=True) + continue + if isinstance(data, list): + logger.info( + "Rerun static entity", + entity_path=entity_path, + archetypes=[type(archetype).__name__ for archetype in data], + ) for archetype in data: rr.log(entity_path, archetype, static=True) else: + logger.info( + "Rerun static entity", + entity_path=entity_path, + archetypes=[type(data).__name__], + ) rr.log(entity_path, data, static=True) @rpc diff --git a/dimos/visualization/rerun/scene_package.py b/dimos/visualization/rerun/scene_package.py new file mode 100644 index 0000000000..a7b29f19ae --- /dev/null +++ b/dimos/visualization/rerun/scene_package.py @@ -0,0 +1,108 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Rerun helpers for cooked scene packages.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass +import math +from pathlib import Path +from typing import Any + +from dimos.simulation.scene_assets.spec import SceneMeshAlignment, ScenePackage, load_scene_package +from dimos.simulation.scenes.catalog import resolve_scene_package + +StaticRerunFactory = Callable[[Any], Any] + + +@dataclass(frozen=True) +class SceneVisualFactory: + visual_path: Path + alignment: SceneMeshAlignment + + def __call__(self, rr: Any) -> list[Any]: + if not self.visual_path.exists(): + return [] + + return [ + rr.Transform3D( + translation=list(self.alignment.translation), + mat3x3=scene_alignment_matrix(self.alignment), + scale=float(self.alignment.scale), + ), + rr.Asset3D( + contents=self.visual_path.read_bytes(), + media_type="model/gltf-binary", + ), + ] + + +def scene_package_static_entities( + scene: str | Path | None, + *, + entity_path: str = "world/scene", +) -> dict[str, StaticRerunFactory]: + """Return static Rerun factories for a cooked scene package visual. + + The scene package visual artifact is stored in source coordinates. The + returned factory logs the package alignment plus the GLB bytes so native and + browser viewers do not need direct filesystem access to the asset. + """ + package = resolve_scene_package_for_rerun(scene) + if package is None or package.visual_path is None or not package.visual_path.exists(): + return {} + + return {entity_path: SceneVisualFactory(package.visual_path, package.alignment)} + + +def resolve_scene_package_for_rerun(scene: str | Path | None) -> ScenePackage | None: + """Resolve normal scene arguments and composed MuJoCo binaries to a package.""" + if scene is None: + return None + + candidate = Path(str(scene).strip()).expanduser() + if candidate.suffix.lower() == ".mjb": + for parent in candidate.resolve().parents: + metadata_path = parent / "scene.meta.json" + if metadata_path.exists(): + return load_scene_package(metadata_path) + return None + + return resolve_scene_package(scene) + + +def scene_alignment_matrix(alignment: SceneMeshAlignment) -> list[list[float]]: + """Return the rotation part of SceneMeshAlignment as a Rerun mat3x3.""" + yaw, pitch, roll = [math.radians(float(v)) for v in alignment.rotation_zyx_deg] + cz, sz = math.cos(yaw), math.sin(yaw) + cy, sy = math.cos(pitch), math.sin(pitch) + cx, sx = math.cos(roll), math.sin(roll) + rz = ((cz, -sz, 0.0), (sz, cz, 0.0), (0.0, 0.0, 1.0)) + ry = ((cy, 0.0, sy), (0.0, 1.0, 0.0), (-sy, 0.0, cy)) + rx = ((1.0, 0.0, 0.0), (0.0, cx, -sx), (0.0, sx, cx)) + + matrix = _matmul(_matmul(rz, ry), rx) + if alignment.y_up: + y_to_z = ((1.0, 0.0, 0.0), (0.0, 0.0, -1.0), (0.0, 1.0, 0.0)) + matrix = _matmul(matrix, y_to_z) + return [list(row) for row in matrix] + + +def _matmul( + a: Sequence[Sequence[float]], + b: Sequence[Sequence[float]], +) -> list[list[float]]: + return [[sum(a[i][k] * b[k][j] for k in range(3)) for j in range(3)] for i in range(3)] diff --git a/dimos/visualization/rerun/test_scene_package.py b/dimos/visualization/rerun/test_scene_package.py new file mode 100644 index 0000000000..afa00fcc5a --- /dev/null +++ b/dimos/visualization/rerun/test_scene_package.py @@ -0,0 +1,108 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path +import pickle +from typing import Any + +import pytest + +from dimos.simulation.scene_assets.spec import SceneMeshAlignment, ScenePackage +from dimos.visualization.rerun.scene_package import ( + resolve_scene_package_for_rerun, + scene_alignment_matrix, + scene_package_static_entities, +) + + +class FakeRerun: + class Transform3D: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + class Asset3D: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + +def test_resolves_composed_mjb_to_scene_package_visual(tmp_path: Path) -> None: + package_dir = tmp_path / "store" + visual_path = package_dir / "browser" / "visual.glb" + visual_path.parent.mkdir(parents=True) + visual_path.write_bytes(b"glb bytes") + composed_path = package_dir / "mujoco" / "composed" / "g1.mjb" + composed_path.parent.mkdir(parents=True) + composed_path.write_bytes(b"mjb bytes") + + alignment = SceneMeshAlignment( + scale=2.0, + translation=(1.0, 2.0, 3.0), + rotation_zyx_deg=(0.0, 0.0, 0.0), + y_up=True, + ) + ScenePackage( + package_dir=package_dir, + source_path=tmp_path / "source.blend", + alignment=alignment, + visual_path=visual_path, + ).write_metadata() + + package = resolve_scene_package_for_rerun(composed_path) + assert package is not None + assert package.visual_path == visual_path.resolve() + + static_entities = scene_package_static_entities(composed_path) + assert set(static_entities) == {"world/scene"} + pickle.dumps(static_entities) + + transform, asset = static_entities["world/scene"](FakeRerun) + assert transform.kwargs == { + "translation": [1.0, 2.0, 3.0], + "mat3x3": [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]], + "scale": 2.0, + } + assert asset.kwargs == { + "contents": b"glb bytes", + "media_type": "model/gltf-binary", + } + + +def test_missing_visual_does_not_add_static_scene(tmp_path: Path) -> None: + package_dir = tmp_path / "store" + composed_path = package_dir / "mujoco" / "composed" / "g1.mjb" + composed_path.parent.mkdir(parents=True) + composed_path.write_bytes(b"mjb bytes") + ScenePackage( + package_dir=package_dir, + source_path=tmp_path / "source.blend", + alignment=SceneMeshAlignment(), + visual_path=None, + ).write_metadata() + + assert scene_package_static_entities(composed_path) == {} + + +def test_scene_alignment_matrix_keeps_z_up_when_requested() -> None: + alignment = SceneMeshAlignment( + rotation_zyx_deg=(90.0, 0.0, 0.0), + y_up=False, + ) + + matrix = scene_alignment_matrix(alignment) + + assert matrix[0][0] == pytest.approx(0.0) + assert matrix[0][1] == pytest.approx(-1.0) + assert matrix[1][0] == pytest.approx(1.0) + assert matrix[1][1] == pytest.approx(0.0) + assert matrix[2] == [0.0, 0.0, 1.0] diff --git a/dimos/visualization/rerun/urdf_robot.py b/dimos/visualization/rerun/urdf_robot.py new file mode 100644 index 0000000000..2476e6192c --- /dev/null +++ b/dimos/visualization/rerun/urdf_robot.py @@ -0,0 +1,217 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Generic Rerun helpers for visualizing URDF robots.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import numpy as np + +from dimos.utils.data import get_data + +JointNameMapper = Callable[[str], str] + + +def default_joint_name_mapper(name: str) -> str: + """Map a DimOS hardware joint name to the common URDF joint name form.""" + short = name.rsplit("/", 1)[-1] + return short if short.endswith("_joint") else f"{short}_joint" + + +def _resolve_urdf_path(path: str | Path) -> Path: + candidate = Path(path).expanduser() + if candidate.is_absolute() or candidate.exists(): + return candidate + return get_data(candidate) + + +def _matrix_to_rerun_transform(rr: Any, matrix: np.ndarray) -> Any: + return rr.Transform3D( + translation=matrix[:3, 3].tolist(), + mat3x3=matrix[:3, :3].tolist(), + ) + + +def _rerun_transform_without_frames(rr: Any, transform: Any) -> Any: + translation = transform.translation.as_arrow_array().to_pylist()[0] + quaternion = transform.quaternion.as_arrow_array().to_pylist()[0] + return rr.Transform3D( + translation=translation, + rotation=rr.Quaternion(xyzw=quaternion), + ) + + +def _mesh_to_rerun(rr: Any, mesh: Any) -> Any: + color = (178, 178, 178, 255) + face_colors = getattr(getattr(mesh, "visual", None), "face_colors", None) + if face_colors is not None and len(face_colors): + rgba = np.asarray(face_colors[0], dtype=np.uint8).tolist() + color = tuple(int(v) for v in rgba) + + return rr.Mesh3D( + vertex_positions=np.asarray(mesh.vertices, dtype=np.float32), + triangle_indices=np.asarray(mesh.faces, dtype=np.uint32), + vertex_normals=np.asarray(mesh.vertex_normals, dtype=np.float32), + albedo_factor=color, + ) + + +def _rerun_path_part(name: str) -> str: + return name.replace("/", "_").replace(" ", "_") + + +def _build_link_paths( + root_path: str, + root_link: str, + joints: list[Any], +) -> dict[str, str]: + link_paths = {root_link: f"{root_path}/{_rerun_path_part(root_link)}"} + remaining = list(joints) + + while remaining: + progressed = False + for joint in remaining[:]: + parent = str(joint.parent) + child = str(joint.child) + parent_path = link_paths.get(parent) + if parent_path is None: + continue + link_paths[child] = f"{parent_path}/{_rerun_path_part(child)}" + remaining.remove(joint) + progressed = True + if not progressed: + for joint in remaining: + link_paths[str(joint.child)] = f"{root_path}/{_rerun_path_part(joint.child)}" + break + + return link_paths + + +@dataclass +class UrdfRobotStaticRerunFactory: + """Log a URDF robot's static visual meshes under a Rerun root path.""" + + urdf_path: str | Path + root_path: str + _robot: Any = field(default=None, init=False, repr=False) + + def __call__(self, rr: Any) -> list[tuple[str, Any]]: + robot = self._load_robot() + link_paths = _build_link_paths( + self.root_path, + str(robot.base_link), + list(robot.robot.joints), + ) + entities: list[tuple[str, Any]] = [ + (self.root_path, rr.Transform3D()), + (link_paths[str(robot.base_link)], rr.Transform3D()), + ] + + geometry_nodes = set(robot.scene.graph.nodes_geometry) + for parent, node, edge_data in robot.scene.graph.to_edgelist(): + if node not in geometry_nodes: + continue + + geometry_name = edge_data.get("geometry") + mesh = robot.scene.geometry.get(geometry_name) + if mesh is None: + continue + + parent_path = link_paths.get( + str(parent), f"{self.root_path}/{_rerun_path_part(parent)}" + ) + path = f"{parent_path}/{_rerun_path_part(node)}" + matrix = np.asarray(edge_data.get("matrix", np.eye(4)), dtype=float) + entities.append((path, _matrix_to_rerun_transform(rr, matrix))) + entities.append((path, _mesh_to_rerun(rr, mesh))) + + return entities + + def _load_robot(self) -> Any: + if self._robot is None: + from yourdfpy import URDF + + self._robot = URDF.load(str(_resolve_urdf_path(self.urdf_path))) + return self._robot + + +@dataclass +class UrdfRobotJointStateRerunFactory: + """Convert JointState-like messages into animated URDF link transforms.""" + + urdf_path: str | Path + root_path: str + joint_name_mapper: JointNameMapper = default_joint_name_mapper + clamp_joint_limits: bool = False + _tree: Any = field(default=None, init=False, repr=False) + _joints: list[Any] = field(default_factory=list, init=False, repr=False) + _joint_paths: dict[str, str] = field(default_factory=dict, init=False, repr=False) + _joint_values: dict[str, float] = field(default_factory=dict, init=False, repr=False) + + def __call__(self, msg: Any) -> list[tuple[str, Any]]: + self._load_tree() + + for name, position in zip(msg.name, msg.position, strict=False): + urdf_joint = self.joint_name_mapper(str(name)) + if urdf_joint in self._joint_values: + self._joint_values[urdf_joint] = float(position) + + import rerun as rr + + return [ + ( + self._joint_paths[joint.name], + _rerun_transform_without_frames( + rr, + joint.compute_transform( + self._joint_values.get(joint.name, 0.0), + clamp=self.clamp_joint_limits, + ), + ), + ) + for joint in self._joints + ] + + def _load_tree(self) -> None: + if self._tree is None: + import rerun.urdf as rr_urdf + from yourdfpy import URDF + + urdf_path = _resolve_urdf_path(self.urdf_path) + robot = URDF.load(str(urdf_path)) + link_paths = _build_link_paths( + self.root_path, + str(robot.base_link), + list(robot.robot.joints), + ) + + self._tree = rr_urdf.UrdfTree.from_file_path(urdf_path) + self._joints = self._tree.joints() + self._joint_paths = { + joint.name: link_paths.get( + joint.child_link, + f"{self.root_path}/{_rerun_path_part(joint.child_link)}", + ) + for joint in self._joints + } + self._joint_values = { + joint.name: 0.0 + for joint in self._joints + if joint.joint_type in {"revolute", "continuous", "prismatic"} + } diff --git a/dimos/web/templates/rerun_dashboard.html b/dimos/web/templates/rerun_dashboard.html index f0792079e3..2cea9fd6be 100644 --- a/dimos/web/templates/rerun_dashboard.html +++ b/dimos/web/templates/rerun_dashboard.html @@ -31,7 +31,7 @@
- +