From c07aed03314e9ee6a8413c7b626579752139cef9 Mon Sep 17 00:00:00 2001 From: Nabla7 Date: Thu, 25 Jun 2026 02:07:02 +0200 Subject: [PATCH 01/10] 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 643f7af3790d23620162348bda9e73465648f9c6 Mon Sep 17 00:00:00 2001 From: Nabla7 Date: Fri, 26 Jun 2026 04:43:26 +0200 Subject: [PATCH 02/10] 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 @@
- +