diff --git a/data/.lfs/g1_urdf.tar.gz b/data/.lfs/g1_urdf.tar.gz
new file mode 100644
index 0000000000..81c7f1db9a
--- /dev/null
+++ b/data/.lfs/g1_urdf.tar.gz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c8e1ed990471461d54069876873b1feddaf23171222413d03ce20a809c1da95b
+size 43232580
diff --git a/data/.lfs/scene_packages.tar.gz b/data/.lfs/scene_packages.tar.gz
new file mode 100644
index 0000000000..a2dad70974
--- /dev/null
+++ b/data/.lfs/scene_packages.tar.gz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c981384e8e6da5f9a0f66d16582337d75e2aacf84b8b18ff185d00f216f983d0
+size 1671132842
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..48edf7e822
--- /dev/null
+++ b/dimos/robot/unitree/g1/assets/g1_29dof.xml
@@ -0,0 +1,300 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
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..90f7f7520a 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,12 +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 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 \\
@@ -43,7 +47,7 @@
from __future__ import annotations
from pathlib import Path
-from typing import Any
+from typing import Any, cast
from dimos.control.components import HardwareComponent, HardwareType
from dimos.control.coordinator import ControlCoordinator, TaskConfig
@@ -59,24 +63,83 @@
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
# 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"
+_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,
@@ -94,43 +157,132 @@
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,
+ **_mujoco_lidar_kwargs(_MUJOCO_LIDAR_CAMERA, _MUJOCO_LIDAR_CAMERAS),
+ 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
+
+ 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)
+ 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}")
+
+ 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,
+ robot_mjcf=_ROBOT_ONLY_MJCF_PATH,
+ robot_meshdir=_ROBOT_MESHDIR,
+ robot_id="",
+ scene_entities=package.entities,
+ headless=True,
+ dof=29,
+ **_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 = 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
_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
@@ -155,6 +307,8 @@
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:
@@ -174,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: dict[str, Any] = {
+ # 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(
@@ -246,8 +403,7 @@ 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),
+ ("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
@@ -258,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(cast("Any", _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 d9889027f4..d8a3c4089b 100644
--- a/dimos/simulation/engines/mujoco_engine.py
+++ b/dimos/simulation/engines/mujoco_engine.py
@@ -18,10 +18,11 @@
from collections.abc import Callable
from dataclasses import dataclass
+import math
from pathlib import Path
import threading
import time
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
import mujoco
import mujoco.viewer as viewer # type: ignore[import-untyped]
@@ -47,6 +48,36 @@
StepHook = Callable[["MujocoEngine"], None]
_MJJNT_FREE = int(mujoco.mjtJoint.mjJNT_FREE) # type: ignore[attr-defined]
+_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 cast("NDArray[np.float64]", directions / norms)
@dataclass
@@ -55,6 +86,8 @@ class CameraConfig:
width: int = 640
height: int = 480
fps: float = 15.0
+ max_geom: int | None = 10000
+ geom_groups: tuple[int, ...] | None = None
@dataclass
@@ -67,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.
@@ -91,29 +154,49 @@ 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,
+ 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
-
- xml_path = self._resolve_xml_path(config_path)
- if assets is not None:
+ self._spawn_xy = spawn_xy
+ self._spawn_z = spawn_z
+ self._spawn_yaw = spawn_yaw
+ self._reset_joint_positions = reset_joint_positions
+
+ 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")
+ 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.
- 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)) # type: ignore[attr-defined]
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)
- self._joint_mappings = build_joint_mappings(self._xml_path, self._model)
+ 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:
self._robot_binding = resolve_robot_sim_binding(
@@ -131,6 +214,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 +227,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
@@ -151,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,
@@ -165,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:
@@ -271,23 +361,77 @@ 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,
+ ) # type: ignore[call-arg]
+ depth_renderer = mujoco.Renderer(
+ self._model,
+ height=cfg.height,
+ width=cfg.width,
+ max_geom=max_geom,
+ ) # type: ignore[call-arg]
depth_renderer.enable_depth_rendering()
+ scene_option = None
+ if cfg.geom_groups is not None:
+ scene_option = mujoco.MjvOption()
+ geomgroup = scene_option.geomgroup # type: ignore[attr-defined]
+ geomgroup[:] = 0
+ for group in cfg.geom_groups:
+ if 0 <= group < len(geomgroup):
+ 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():
@@ -295,10 +439,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(
@@ -312,21 +460,140 @@ 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( # type: ignore[attr-defined]
+ 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():
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) # type: ignore[attr-defined]
+ 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
# 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()
+ 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)
@@ -343,6 +610,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
@@ -486,6 +754,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( # type: ignore[attr-defined]
+ 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:
@@ -519,9 +871,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 ffc2effa39..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
@@ -188,6 +190,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 +207,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
@@ -218,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;
@@ -294,6 +329,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 +369,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
@@ -349,36 +389,81 @@ 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.
- self._engine = MujocoEngine(
- config_path=Path(self.config.address),
+ engine_kwargs: dict[str, Any] = dict(
headless=self.config.headless,
cameras=cameras,
- assets=engine_assets,
+ raycast_lidars=raycast_lidars,
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 +508,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(
@@ -459,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(
@@ -472,11 +562,67 @@ 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, spawn_penetrators
+
+ if self.config.robot_mjcf is None:
+ raise RuntimeError("MujocoSimModule: robot_mjcf is required for composition")
+
+ 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)
+
+ 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:
+ 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:
self._stop_event.set()
@@ -509,6 +655,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.
@@ -737,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
@@ -766,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 1a37941ef0..f0fb1eb25c 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,232 @@ 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,
+ },
+ }
+
+
+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")
+ 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_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"
+ _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/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
new file mode 100644
index 0000000000..fa971d9d1f
--- /dev/null
+++ b/dimos/simulation/mujoco/entity_scene.py
@@ -0,0 +1,361 @@
+# 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 _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"]
+
+
+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_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."""
+ 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
+
+ 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")
+ 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):
+ 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)
+ 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 ""
+ entity_id = _entity_id_from_geom_name(name)
+ if entity_id is not None:
+ bad.add(entity_id)
+ 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
diff --git a/dimos/simulation/scene_assets/spec.py b/dimos/simulation/scene_assets/spec.py
new file mode 100644
index 0000000000..dfcf58b92a
--- /dev/null
+++ b/dimos/simulation/scene_assets/spec.py
@@ -0,0 +1,313 @@
+# Copyright 2025-2026 Dimensional Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""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_visuals: dict[str, Path] = field(default_factory=dict)
+ 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_visuals": _serialize_browser_visuals(
+ self.browser_visuals,
+ 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 browser_visual_path(self, target: str) -> Path | None:
+ """Return a visual GLB for a specific browser/viewer backend.
+
+ Scene packages can carry multiple browser-facing visuals because viewer
+ support for glTF extensions differs. For example, Rerun currently needs
+ conservative GLBs, while Babylon can use more web-oriented output.
+
+ Older packages only have ``browser_visual``. Rerun may use that legacy
+ artifact; other targets should cook their own named asset.
+ """
+ target_key = target.strip().lower()
+ if target_key in self.browser_visuals:
+ return self.browser_visuals[target_key]
+ if target_key == "rerun":
+ return self.browser_visuals.get("default") or self.visual_path
+ return self.browser_visuals.get("default")
+
+ 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)
+ browser_visuals = _resolve_browser_visuals(artifacts.get("browser_visuals"), package_dir)
+ legacy_visual_path = _resolve_package_path(artifacts.get("browser_visual"), package_dir)
+ return ScenePackage(
+ package_dir=package_dir,
+ source_path=Path(raw["source_path"]),
+ alignment=align,
+ visual_path=legacy_visual_path,
+ browser_visuals=browser_visuals,
+ 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()
+
+
+def _serialize_browser_visuals(
+ visuals: dict[str, Path],
+ package_dir: Path,
+) -> dict[str, str]:
+ return {
+ str(target).strip().lower(): serialized
+ for target, path in sorted(visuals.items())
+ if (serialized := _serialize_package_path(path, package_dir)) is not None
+ }
+
+
+def _resolve_browser_visuals(raw: Any, package_dir: Path) -> dict[str, Path]:
+ if not isinstance(raw, dict):
+ return {}
+ resolved: dict[str, Path] = {}
+ for target, path in raw.items():
+ if not isinstance(target, str) or not isinstance(path, str):
+ continue
+ visual_path = _resolve_package_path(path, package_dir)
+ if visual_path is not None:
+ resolved[target.strip().lower()] = visual_path
+ return resolved
+
+
+_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)
diff --git a/dimos/simulation/scenes/README.md b/dimos/simulation/scenes/README.md
new file mode 100644
index 0000000000..4f5b5d5b13
--- /dev/null
+++ b/dimos/simulation/scenes/README.md
@@ -0,0 +1,154 @@
+# Scene Packages
+
+A scene package is a cooked environment that runtime code can load without
+knowing how the environment was produced.
+
+Before scene packages, a scene was usually a loose set of files: MuJoCo XML,
+mesh files, browser assets, object metadata, and local path assumptions. A
+consumer had to know the folder layout for that one scene. Scene packages make
+that explicit: load `scene.meta.json` once, then ask the package for the
+artifact needed by the simulator, viewer, or planner.
+
+Scene packages are runtime artifacts. Cooking tools may use Blender, GLB
+optimizers, collision decomposition, or manual sidecars, but runtime modules
+should not depend on those details.
+
+## What A Package Can Contain
+
+A package can carry several assets for different consumers:
+
+- MuJoCo XML for physics simulation
+- optional MuJoCo `.mjb` binaries for faster loading of large scenes
+- Rerun or browser visual assets
+- future visual assets such as Gaussian splats
+- static environment geometry
+- dynamic or separately spawned entities
+- object prototypes reused by many instances
+- coordinate-frame and scale alignment metadata
+
+Not every package needs every artifact. A viewer can load the visual asset, a
+simulator can load the physics asset, and another renderer can ask for its own
+target-specific output later.
+
+## Runtime API
+
+Use `resolve_scene_package()` for the same values accepted by `--scene`:
+
+```python
+from dimos.simulation.scenes.catalog import resolve_scene_package
+
+package = resolve_scene_package(global_config.scene)
+```
+
+Supported inputs:
+
+- `None` or `--scene none`
+- named packages such as `office` or `supermarket`
+- path to `scene.meta.json`
+- path to a scene package directory
+- path to a precompiled MuJoCo `.mjb` in consumers that support direct binary loading
+
+Use `load_scene_package()` when you already have the exact metadata path:
+
+```python
+from dimos.simulation.scene_assets.spec import load_scene_package
+
+package = load_scene_package("data/scene_packages/dimos_office/scene.meta.json")
+```
+
+Once resolved, consumers should use the package object instead of hardcoded
+scene-specific paths:
+
+```python
+if package is not None:
+ mujoco_xml = package.mujoco_scene_path
+ rerun_visual = package.browser_visual_path("rerun")
+ browser_collision = package.browser_collision_path
+ objects_json = package.objects_path
+ entities = package.entities
+ alignment = package.alignment
+```
+
+`alignment` describes how source assets are scaled and oriented relative to the
+runtime world. This lets viewers and simulators agree on frame conventions.
+
+Browser-facing assets are target-specific. Rerun should ask for
+`browser_visual_path("rerun")`; a browser renderer such as Babylon/PimSim can
+ask for `browser_visual_path("babylon")`. If a target returns `None`, the
+package was not cooked for that target.
+
+## Named Packages
+
+The current named packages are:
+
+```text
+none no cooked scene package
+office cooked DimOS office package
+supermarket cooked supermarket package
+```
+
+## MuJoCo Loading
+
+Normal packages are robot-agnostic:
+
+```text
+scene.meta.json -> wrapper.xml + entities -> attach robot MJCF -> compile MjModel
+```
+
+Large scenes may also ship precompiled MuJoCo binaries:
+
+```text
+wrapper.xml + robot MJCF + spawn/entity selection -> composed/.mjb
+```
+
+Loading a composed `.mjb` is faster, but that file is specific to one robot,
+spawn pose, entity selection, and scene revision. Keep the scene package
+metadata as the source of truth and treat composed binaries as cache artifacts.
+
+## Run G1 With A Scene
+
+No scene:
+
+```bash
+dimos \
+ --simulation mujoco \
+ --scene none \
+ --viewer rerun \
+ --rerun-open native \
+ --n-workers 12 \
+ run unitree-g1-groot-wbc
+```
+
+Office:
+
+```bash
+dimos \
+ --simulation mujoco \
+ --scene office \
+ --viewer rerun \
+ --rerun-open native \
+ --n-workers 12 \
+ run unitree-g1-groot-wbc
+```
+
+Supermarket:
+
+```bash
+dimos \
+ --simulation mujoco \
+ --scene supermarket \
+ --viewer rerun \
+ --rerun-open native \
+ --n-workers 12 \
+ run unitree-g1-groot-wbc
+```
+
+Prefer headless MuJoCo with Rerun native for normal testing. Opening the MuJoCo
+viewer with `mujocosimmodule.headless=false` is useful for contact debugging,
+but it can run much slower.
+
+## Boundaries
+
+This package contract is for runtime loading. Cooking logic belongs in the
+scene cooking tools. Runtime modules should stay on the `ScenePackage` API and
+avoid depending on Blender, source-scene naming, or cooking heuristics.
diff --git a/dimos/simulation/scenes/catalog.py b/dimos/simulation/scenes/catalog.py
new file mode 100644
index 0000000000..be3edb77c3
--- /dev/null
+++ b/dimos/simulation/scenes/catalog.py
@@ -0,0 +1,112 @@
+# 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 = "office"
+_DISABLED_SCENE_NAMES = {"", "none", "off", "disabled", "false", "0"}
+_ALIASES = {
+ "office": DEFAULT_SCENE,
+ "dimos-office": DEFAULT_SCENE,
+ "dimos_office": DEFAULT_SCENE,
+ "supermarket": "supermarket",
+}
+_PACKAGE_DIRS = {
+ DEFAULT_SCENE: "dimos_office",
+ "supermarket": "supermarket",
+}
+
+
+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:
+ return None
+
+ 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)
+ rerun_visual_path = package.browser_visual_path("rerun")
+ if (
+ rerun_visual_path is not None
+ and package.browser_collision_path is not None
+ and package.mujoco_scene_path is not None
+ and rerun_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")
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/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..185159ca0a
--- /dev/null
+++ b/dimos/visualization/rerun/scene_package.py
@@ -0,0 +1,112 @@
+# 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:
+ return {}
+
+ visual_path = package.browser_visual_path("rerun")
+ if visual_path is None or not visual_path.exists():
+ return {}
+
+ return {entity_path: SceneVisualFactory(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..b5a5f185f2
--- /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.rerun.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,
+ browser_visuals={"rerun": visual_path},
+ ).write_metadata()
+
+ package = resolve_scene_package_for_rerun(composed_path)
+ assert package is not None
+ assert package.browser_visual_path("rerun") == 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..2e7e7293b4
--- /dev/null
+++ b/dimos/visualization/rerun/urdf_robot.py
@@ -0,0 +1,218 @@
+# 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: tuple[int, int, int, int] = (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()
+ if len(rgba) >= 4:
+ color = (int(rgba[0]), int(rgba[1]), int(rgba[2]), int(rgba[3]))
+
+ 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 # type: ignore[import-not-found]
+
+ 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 # type: ignore[import-not-found]
+
+ 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/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