From ee68567f67f3e6bf5566467ae529a6966d197ed9 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 9 Feb 2026 16:08:08 -0800 Subject: [PATCH 1/6] with seperate preview urdf --- .../planning/world/drake_world.py | 133 +++++++++++++++--- 1 file changed, 116 insertions(+), 17 deletions(-) diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index 2aad2fc163..3016ec8f82 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -17,7 +17,7 @@ from __future__ import annotations from contextlib import contextmanager -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from threading import RLock from typing import TYPE_CHECKING, Any @@ -52,6 +52,7 @@ Cylinder, GeometryInstance, GeometrySet, + IllustrationProperties, MakePhongIllustrationProperties, Meshcat, MeshcatVisualizer, @@ -59,6 +60,7 @@ ProximityProperties, Rgba, Role, + RoleAssign, SceneGraph, Sphere, ) @@ -89,6 +91,8 @@ class _RobotData: joint_indices: list[int] # Indices into plant's position vector ee_frame: Any # BodyFrame for end-effector base_frame: Any # BodyFrame for base + preview_model_instance: Any = None # ModelInstanceIndex for preview (yellow) robot + preview_joint_indices: list[int] = field(default_factory=list) @dataclass @@ -167,6 +171,12 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: ).body_frame() base_frame = self._plant.GetBodyByName(config.base_link, model_instance).body_frame() + # Load a second copy of the URDF as the preview (yellow ghost) robot + preview_model_instance = None + if self._enable_viz: + preview_model_instance = self._load_urdf(config) + self._weld_base_if_needed(config, preview_model_instance) + self._robots[robot_id] = _RobotData( robot_id=robot_id, config=config, @@ -174,6 +184,7 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: joint_indices=[], ee_frame=ee_frame, base_frame=base_frame, + preview_model_instance=preview_model_instance, ) logger.info(f"Added robot '{robot_id}' ({config.name})") @@ -479,6 +490,35 @@ def clear_obstacles(self) -> None: for obs_id in obstacle_ids: self.remove_obstacle(obs_id) + # ============= Preview Robot Setup ============= + + def _set_preview_colors(self) -> None: + """Set all preview robot visual geometries to yellow/semi-transparent.""" + source_id = self._plant.get_source_id() + preview_color = Rgba(1.0, 0.8, 0.0, 0.4) + + for robot_data in self._robots.values(): + if robot_data.preview_model_instance is None: + continue + for body_idx in self._plant.GetBodyIndices(robot_data.preview_model_instance): + body = self._plant.get_body(body_idx) + for geom_id in self._plant.GetVisualGeometriesForBody(body): + props = IllustrationProperties() + props.AddProperty("phong", "diffuse", preview_color) + self._scene_graph.AssignRole(source_id, geom_id, props, RoleAssign.kReplace) + + def _remove_preview_collision_roles(self) -> None: + """Remove proximity (collision) role from all preview robot geometries.""" + source_id = self._plant.get_source_id() + + for robot_data in self._robots.values(): + if robot_data.preview_model_instance is None: + continue + for body_idx in self._plant.GetBodyIndices(robot_data.preview_model_instance): + body = self._plant.get_body(body_idx) + for geom_id in self._plant.GetCollisionGeometriesForBody(body): + self._scene_graph.RemoveRole(source_id, geom_id, Role.kProximity) + # ============= Lifecycle ============= def finalize(self) -> None: @@ -491,7 +531,7 @@ def finalize(self) -> None: # Finalize plant self._plant.Finalize() - # Compute joint indices for each robot + # Compute joint indices for each robot (live + preview) for robot_id, robot_data in self._robots.items(): joint_indices: list[int] = [] for joint_name in robot_data.config.joint_names: @@ -502,9 +542,28 @@ def finalize(self) -> None: robot_data.joint_indices = joint_indices logger.debug(f"Robot '{robot_id}' joint indices: {joint_indices}") + # Compute preview joint indices + if robot_data.preview_model_instance is not None: + preview_indices: list[int] = [] + for joint_name in robot_data.config.joint_names: + joint = self._plant.GetJointByName( + joint_name, robot_data.preview_model_instance + ) + start_idx = joint.position_start() + num_positions = joint.num_positions() + preview_indices.extend(range(start_idx, start_idx + num_positions)) + robot_data.preview_joint_indices = preview_indices + logger.debug(f"Robot '{robot_id}' preview joint indices: {preview_indices}") + # Setup collision filters self._setup_collision_filters() + # Remove collision roles from preview robots (visual-only) + self._remove_preview_collision_roles() + + # Set preview robots to yellow/semi-transparent + self._set_preview_colors() + # Register obstacle source for dynamic obstacles self._obstacle_source_id = self._scene_graph.RegisterSource("dynamic_obstacles") @@ -539,6 +598,9 @@ def finalize(self) -> None: self._meshcat_visualizer.ForcedPublish( self._diagram.GetSubsystemContext(self._meshcat_visualizer, self._live_context) ) + # Hide all preview robots initially + for robot_id in self._robots: + self.hide_preview(robot_id) @property def is_finalized(self) -> bool: @@ -864,13 +926,48 @@ def publish_visualization(self, ctx: Context | None = None) -> None: self._diagram.GetSubsystemContext(self._meshcat_visualizer, ctx) ) + def _set_preview_positions( + self, plant_ctx: Context, robot_id: WorldRobotID, positions: NDArray[np.float64] + ) -> None: + """Set preview robot positions in a plant context.""" + robot_data = self._robots.get(robot_id) + if robot_data is None or robot_data.preview_model_instance is None: + return + + full_positions = self._plant.GetPositions(plant_ctx).copy() + for i, idx in enumerate(robot_data.preview_joint_indices): + full_positions[idx] = positions[i] + self._plant.SetPositions(plant_ctx, full_positions) + + def show_preview(self, robot_id: WorldRobotID) -> None: + """Show the preview (yellow ghost) robot in Meshcat.""" + if self._meshcat is None: + return + robot_data = self._robots.get(robot_id) + if robot_data is None or robot_data.preview_model_instance is None: + return + model_name = self._plant.GetModelInstanceName(robot_data.preview_model_instance) + self._meshcat.SetProperty(f"visualizer/{model_name}", "visible", True) + + def hide_preview(self, robot_id: WorldRobotID) -> None: + """Hide the preview (yellow ghost) robot in Meshcat.""" + if self._meshcat is None: + return + robot_data = self._robots.get(robot_id) + if robot_data is None or robot_data.preview_model_instance is None: + return + model_name = self._plant.GetModelInstanceName(robot_data.preview_model_instance) + self._meshcat.SetProperty(f"visualizer/{model_name}", "visible", False) + def animate_path( self, robot_id: WorldRobotID, path: JointPath, duration: float = 3.0, ) -> None: - """Animate a path in Meshcat visualization. + """Animate a path using the preview (yellow ghost) robot. + + The live robot stays in place while the preview robot shows the planned path. Args: robot_id: Robot to animate @@ -882,22 +979,24 @@ def animate_path( if self._meshcat is None or len(path) < 2: return - # Capture current states of all OTHER robots so they don't snap to zero - other_robot_states: dict[WorldRobotID, JointState] = {} - for rid, _robot_data in self._robots.items(): - if rid != robot_id: - other_robot_states[rid] = self.get_joint_state(self.get_live_context(), rid) + robot_data = self._robots.get(robot_id) + if robot_data is None or robot_data.preview_model_instance is None: + return + + self.show_preview(robot_id) dt = duration / (len(path) - 1) - for joint_state in path: - with self.scratch_context() as ctx: - # Restore other robots to their current states - for rid, state in other_robot_states.items(): - self.set_joint_state(ctx, rid, state) - # Set animated robot's state - self.set_joint_state(ctx, robot_id, joint_state) - self.publish_visualization(ctx) - time.sleep(dt) + try: + for joint_state in path: + positions = np.array(joint_state.position, dtype=np.float64) + with self._lock: + self._set_preview_positions(self._plant_context, robot_id, positions) + # Publish live context — renders both live (current) and preview (animated) + self.publish_visualization() + time.sleep(dt) + finally: + self.hide_preview(robot_id) + self.publish_visualization() # ============= Direct Access (use with caution) ============= From 0eb9a0552a85edfa0032f1dbf3d8cf1a6a32bcf3 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 9 Feb 2026 21:13:25 -0800 Subject: [PATCH 2/6] running meshcat on its dedicated thread allows for real time preview update --- dimos/manipulation/manipulation_module.py | 2 +- .../planning/world/drake_world.py | 80 ++++++++++--------- 2 files changed, 45 insertions(+), 37 deletions(-) diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index c57663f775..8401e22ffa 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -492,7 +492,7 @@ def preview_path(self, duration: float = 3.0, robot_name: RobotName | None = Non return False # Interpolate and animate - interpolated = interpolate_path(planned_path, resolution=0.02) + interpolated = interpolate_path(planned_path, resolution=0.1) self._world_monitor.world.animate_path(robot_id, interpolated, duration) return True diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index 3016ec8f82..aee4812770 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -16,10 +16,11 @@ from __future__ import annotations +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass, field from pathlib import Path -from threading import RLock +from threading import RLock, current_thread from typing import TYPE_CHECKING, Any import numpy as np @@ -128,11 +129,17 @@ def __init__(self, time_step: float = 0.0, enable_viz: bool = False): # with the same URDF (e.g., 4 XArm6 arms all have model name "UF_ROBOT") self._parser.SetAutoRenaming(True) - # Visualization + # Visualization — all Meshcat calls must happen on the thread that created + # the Meshcat instance (Drake enforces thread affinity via SystemExit). + # A single-thread executor ensures every viz call runs on that thread. + self._viz_executor: ThreadPoolExecutor | None = None + self._viz_thread: Any = None # Thread object for re-entrancy detection self._meshcat: Meshcat | None = None self._meshcat_visualizer: MeshcatVisualizer | None = None if enable_viz: - self._meshcat = Meshcat() + self._viz_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="meshcat") + self._viz_thread = self._viz_executor.submit(current_thread).result() + self._meshcat = self._viz_executor.submit(Meshcat).result() # Create model instance for obstacles self._obstacles_model_instance = self._plant.AddModelInstance("obstacles") @@ -413,8 +420,8 @@ def _add_obstacle_to_meshcat(self, obstacle: Obstacle, obstacle_id: str) -> None # Create Drake shape and add to Meshcat drake_shape = self._create_shape(obstacle) - self._meshcat.SetObject(path, drake_shape, rgba) - self._meshcat.SetTransform(path, transform) + self._on_viz_thread(self._meshcat.SetObject, path, drake_shape, rgba) + self._on_viz_thread(self._meshcat.SetTransform, path, transform) def _pose_to_rigid_transform(self, pose: PoseStamped) -> Any: """Convert PoseStamped to Drake RigidTransform.""" @@ -456,7 +463,7 @@ def remove_obstacle(self, obstacle_id: str) -> bool: # Also remove from Meshcat if self._meshcat is not None: path = f"obstacles/{obstacle_id}" - self._meshcat.Delete(path) + self._on_viz_thread(self._meshcat.Delete, path) del self._obstacles[obstacle_id] logger.debug(f"Removed obstacle '{obstacle_id}'") @@ -475,7 +482,7 @@ def update_obstacle_pose(self, obstacle_id: str, pose: PoseStamped) -> bool: if self._meshcat is not None: path = f"obstacles/{obstacle_id}" transform = self._pose_to_rigid_transform(pose) - self._meshcat.SetTransform(path, transform) + self._on_viz_thread(self._meshcat.SetTransform, path, transform) # Note: SceneGraph geometry pose is fixed after registration # Meshcat is updated for visualization, but collision checking @@ -593,11 +600,9 @@ def finalize(self) -> None: self._finalized = True logger.info(f"World finalized with {len(self._robots)} robots") - # Initial visualization publish + # Initial visualization publish (routed to Meshcat thread) if self._meshcat_visualizer is not None: - self._meshcat_visualizer.ForcedPublish( - self._diagram.GetSubsystemContext(self._meshcat_visualizer, self._live_context) - ) + self.publish_visualization() # Hide all preview robots initially for robot_id in self._robots: self.hide_preview(robot_id) @@ -902,10 +907,16 @@ def get_jacobian(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.float # ============= Visualization ============= + def _on_viz_thread(self, fn: Any, *args: Any, **kwargs: Any) -> Any: + """Run fn on the Meshcat creator thread. Re-entrant safe.""" + if self._viz_executor is None or current_thread() is self._viz_thread: + return fn(*args, **kwargs) + return self._viz_executor.submit(fn, *args, **kwargs).result() + def get_visualization_url(self) -> str | None: """Get visualization URL if enabled.""" if self._meshcat is not None: - url: str = self._meshcat.web_url() + url: str = self._on_viz_thread(self._meshcat.web_url) return url return None @@ -922,9 +933,8 @@ def publish_visualization(self, ctx: Context | None = None) -> None: ctx = self._live_context if ctx is not None: - self._meshcat_visualizer.ForcedPublish( - self._diagram.GetSubsystemContext(self._meshcat_visualizer, ctx) - ) + viz_ctx = self._diagram.GetSubsystemContext(self._meshcat_visualizer, ctx) + self._on_viz_thread(self._meshcat_visualizer.ForcedPublish, viz_ctx) def _set_preview_positions( self, plant_ctx: Context, robot_id: WorldRobotID, positions: NDArray[np.float64] @@ -947,7 +957,7 @@ def show_preview(self, robot_id: WorldRobotID) -> None: if robot_data is None or robot_data.preview_model_instance is None: return model_name = self._plant.GetModelInstanceName(robot_data.preview_model_instance) - self._meshcat.SetProperty(f"visualizer/{model_name}", "visible", True) + self._on_viz_thread(self._meshcat.SetProperty, f"visualizer/{model_name}", "visible", True) def hide_preview(self, robot_id: WorldRobotID) -> None: """Hide the preview (yellow ghost) robot in Meshcat.""" @@ -957,7 +967,7 @@ def hide_preview(self, robot_id: WorldRobotID) -> None: if robot_data is None or robot_data.preview_model_instance is None: return model_name = self._plant.GetModelInstanceName(robot_data.preview_model_instance) - self._meshcat.SetProperty(f"visualizer/{model_name}", "visible", False) + self._on_viz_thread(self._meshcat.SetProperty, f"visualizer/{model_name}", "visible", False) def animate_path( self, @@ -968,14 +978,8 @@ def animate_path( """Animate a path using the preview (yellow ghost) robot. The live robot stays in place while the preview robot shows the planned path. - - Args: - robot_id: Robot to animate - path: List of joint states forming the path - duration: Total animation duration in seconds + The entire animation runs on the Meshcat thread to satisfy Drake's thread affinity. """ - import time - if self._meshcat is None or len(path) < 2: return @@ -983,20 +987,24 @@ def animate_path( if robot_data is None or robot_data.preview_model_instance is None: return - self.show_preview(robot_id) + def _do_animate() -> None: + import time - dt = duration / (len(path) - 1) - try: - for joint_state in path: - positions = np.array(joint_state.position, dtype=np.float64) - with self._lock: - self._set_preview_positions(self._plant_context, robot_id, positions) - # Publish live context — renders both live (current) and preview (animated) + self.show_preview(robot_id) + dt = duration / (len(path) - 1) + try: + for joint_state in path: + positions = np.array(joint_state.position, dtype=np.float64) + with self._lock: + self._set_preview_positions(self._plant_context, robot_id, positions) + self.publish_visualization() + time.sleep(dt) + finally: + self.hide_preview(robot_id) self.publish_visualization() - time.sleep(dt) - finally: - self.hide_preview(robot_id) - self.publish_visualization() + + # Submit to viz thread so ForcedPublish runs on the Meshcat creator thread + self._on_viz_thread(_do_animate) # ============= Direct Access (use with caution) ============= From c7a984bdf4427745e28ce6020ee23a2f342f12c7 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 9 Feb 2026 21:40:28 -0800 Subject: [PATCH 3/6] added meshcat viz executor shutdown --- dimos/manipulation/planning/monitor/world_monitor.py | 2 ++ dimos/manipulation/planning/spec/protocols.py | 4 ++++ dimos/manipulation/planning/world/drake_world.py | 7 +++++++ 3 files changed, 13 insertions(+) diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index 30c1611e54..7c27753658 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -178,6 +178,8 @@ def stop_all_monitors(self) -> None: logger.info("All monitors stopped") + self._world.close() + # ============= Message Handlers ============= def on_joint_state(self, msg: JointState, robot_id: WorldRobotID | None = None) -> None: diff --git a/dimos/manipulation/planning/spec/protocols.py b/dimos/manipulation/planning/spec/protocols.py index 9f4cf27bf4..dea4718abb 100644 --- a/dimos/manipulation/planning/spec/protocols.py +++ b/dimos/manipulation/planning/spec/protocols.py @@ -178,6 +178,10 @@ def animate_path(self, robot_id: WorldRobotID, path: JointPath, duration: float """Animate a path in visualization.""" ... + def close(self) -> None: + """Release visualization resources.""" + ... + @runtime_checkable class KinematicsSpec(Protocol): diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index aee4812770..1729905d8c 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -996,6 +996,7 @@ def _do_animate() -> None: for joint_state in path: positions = np.array(joint_state.position, dtype=np.float64) with self._lock: + assert self._plant_context is not None self._set_preview_positions(self._plant_context, robot_id, positions) self.publish_visualization() time.sleep(dt) @@ -1006,6 +1007,12 @@ def _do_animate() -> None: # Submit to viz thread so ForcedPublish runs on the Meshcat creator thread self._on_viz_thread(_do_animate) + def close(self) -> None: + """Shut down the viz executor thread.""" + if self._viz_executor is not None: + self._viz_executor.shutdown(wait=False) + self._viz_executor = None + # ============= Direct Access (use with caution) ============= @property From dfd2233c7d28bc6de171770dc01997aa15df2f51 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Tue, 10 Feb 2026 11:53:57 -0800 Subject: [PATCH 4/6] removed sleep on meshcat thread now time.sleep is only called in rpc thread --- .../planning/world/drake_world.py | 33 +++++++------------ 1 file changed, 12 insertions(+), 21 deletions(-) diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index 1729905d8c..a8ea51b3c2 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -977,8 +977,7 @@ def animate_path( ) -> None: """Animate a path using the preview (yellow ghost) robot. - The live robot stays in place while the preview robot shows the planned path. - The entire animation runs on the Meshcat thread to satisfy Drake's thread affinity. + The preview stays visible after animation completes. """ if self._meshcat is None or len(path) < 2: return @@ -987,25 +986,17 @@ def animate_path( if robot_data is None or robot_data.preview_model_instance is None: return - def _do_animate() -> None: - import time - - self.show_preview(robot_id) - dt = duration / (len(path) - 1) - try: - for joint_state in path: - positions = np.array(joint_state.position, dtype=np.float64) - with self._lock: - assert self._plant_context is not None - self._set_preview_positions(self._plant_context, robot_id, positions) - self.publish_visualization() - time.sleep(dt) - finally: - self.hide_preview(robot_id) - self.publish_visualization() - - # Submit to viz thread so ForcedPublish runs on the Meshcat creator thread - self._on_viz_thread(_do_animate) + import time + + self.show_preview(robot_id) + dt = duration / (len(path) - 1) + for joint_state in path: + positions = np.array(joint_state.position, dtype=np.float64) + with self._lock: + assert self._plant_context is not None + self._set_preview_positions(self._plant_context, robot_id, positions) + self.publish_visualization() + time.sleep(dt) def close(self) -> None: """Shut down the viz executor thread.""" From a685eeb81be5d77792fac904468d78851c26e309 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Tue, 10 Feb 2026 11:54:25 -0800 Subject: [PATCH 5/6] preview urdf is now persistent and does not disappear after preview --- dimos/manipulation/manipulation_module.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 8401e22ffa..33dea33697 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -382,6 +382,15 @@ def _fail(self, msg: str) -> bool: self._error_message = msg return False + def _dismiss_preview(self, robot_id: WorldRobotID) -> None: + """Hide the preview ghost if the world supports it.""" + if self._world_monitor is None: + return + world = self._world_monitor.world + if hasattr(world, "hide_preview"): + world.hide_preview(robot_id) # type: ignore[attr-defined] + world.publish_visualization() + @rpc def plan_to_pose(self, pose: Pose, robot_name: RobotName | None = None) -> bool: """Plan motion to pose. Use preview_path() then execute(). @@ -442,6 +451,7 @@ def _plan_path_only( ) -> bool: """Plan path from current position to goal, store result.""" assert self._world_monitor and self._planner # guaranteed by _begin_planning + self._dismiss_preview(robot_id) start = self._world_monitor.get_current_joint_state(robot_id) if start is None: return self._fail("No joint state") From 8a48937a27ab97b4b0b7e8510b016b28fc688c2b Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Tue, 10 Feb 2026 12:01:18 -0800 Subject: [PATCH 6/6] wrapped meshcat threadexecutor in its class --- .../planning/world/drake_world.py | 98 ++++++++++++------- 1 file changed, 61 insertions(+), 37 deletions(-) diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index a8ea51b3c2..532ca6d548 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -106,6 +106,50 @@ class _ObstacleData: source_id: Any # SourceId +class _ThreadSafeMeshcat: + """Wraps Drake Meshcat so all calls run on the creator thread. + + Drake throws SystemExit from non-creator threads for every Meshcat operation. + This class creates a single-thread executor, constructs Meshcat on it, + and proxies all calls through it. + """ + + def __init__(self) -> None: + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="meshcat") + self._thread = self._executor.submit(current_thread).result() + self._inner: Meshcat = self._executor.submit(Meshcat).result() + + def _call(self, fn: Any, *args: Any, **kwargs: Any) -> Any: + if current_thread() is self._thread: + return fn(*args, **kwargs) + return self._executor.submit(fn, *args, **kwargs).result() + + # --- Meshcat proxies --- + + def SetObject(self, *args: Any, **kwargs: Any) -> Any: + return self._call(self._inner.SetObject, *args, **kwargs) + + def SetTransform(self, *args: Any, **kwargs: Any) -> Any: + return self._call(self._inner.SetTransform, *args, **kwargs) + + def SetProperty(self, *args: Any, **kwargs: Any) -> Any: + return self._call(self._inner.SetProperty, *args, **kwargs) + + def Delete(self, *args: Any, **kwargs: Any) -> Any: + return self._call(self._inner.Delete, *args, **kwargs) + + def web_url(self) -> str: + result: str = self._call(self._inner.web_url) + return result + + def forced_publish(self, visualizer: Any, viz_ctx: Any) -> None: + """Run MeshcatVisualizer.ForcedPublish on the creator thread.""" + self._call(visualizer.ForcedPublish, viz_ctx) + + def close(self) -> None: + self._executor.shutdown(wait=False) + + class DrakeWorld(WorldSpec): """Drake implementation of WorldSpec with MultibodyPlant, SceneGraph, optional Meshcat.""" @@ -129,17 +173,11 @@ def __init__(self, time_step: float = 0.0, enable_viz: bool = False): # with the same URDF (e.g., 4 XArm6 arms all have model name "UF_ROBOT") self._parser.SetAutoRenaming(True) - # Visualization — all Meshcat calls must happen on the thread that created - # the Meshcat instance (Drake enforces thread affinity via SystemExit). - # A single-thread executor ensures every viz call runs on that thread. - self._viz_executor: ThreadPoolExecutor | None = None - self._viz_thread: Any = None # Thread object for re-entrancy detection - self._meshcat: Meshcat | None = None + # Visualization — wrapped to enforce Drake's thread affinity + self._meshcat: _ThreadSafeMeshcat | None = None self._meshcat_visualizer: MeshcatVisualizer | None = None if enable_viz: - self._viz_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="meshcat") - self._viz_thread = self._viz_executor.submit(current_thread).result() - self._meshcat = self._viz_executor.submit(Meshcat).result() + self._meshcat = _ThreadSafeMeshcat() # Create model instance for obstacles self._obstacles_model_instance = self._plant.AddModelInstance("obstacles") @@ -420,8 +458,8 @@ def _add_obstacle_to_meshcat(self, obstacle: Obstacle, obstacle_id: str) -> None # Create Drake shape and add to Meshcat drake_shape = self._create_shape(obstacle) - self._on_viz_thread(self._meshcat.SetObject, path, drake_shape, rgba) - self._on_viz_thread(self._meshcat.SetTransform, path, transform) + self._meshcat.SetObject(path, drake_shape, rgba) + self._meshcat.SetTransform(path, transform) def _pose_to_rigid_transform(self, pose: PoseStamped) -> Any: """Convert PoseStamped to Drake RigidTransform.""" @@ -463,7 +501,7 @@ def remove_obstacle(self, obstacle_id: str) -> bool: # Also remove from Meshcat if self._meshcat is not None: path = f"obstacles/{obstacle_id}" - self._on_viz_thread(self._meshcat.Delete, path) + self._meshcat.Delete(path) del self._obstacles[obstacle_id] logger.debug(f"Removed obstacle '{obstacle_id}'") @@ -482,7 +520,7 @@ def update_obstacle_pose(self, obstacle_id: str, pose: PoseStamped) -> bool: if self._meshcat is not None: path = f"obstacles/{obstacle_id}" transform = self._pose_to_rigid_transform(pose) - self._on_viz_thread(self._meshcat.SetTransform, path, transform) + self._meshcat.SetTransform(path, transform) # Note: SceneGraph geometry pose is fixed after registration # Meshcat is updated for visualization, but collision checking @@ -581,7 +619,7 @@ def finalize(self) -> None: self._meshcat_visualizer = MeshcatVisualizer.AddToBuilder( self._builder, self._scene_graph, - self._meshcat, + self._meshcat._inner, params, ) @@ -907,34 +945,21 @@ def get_jacobian(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.float # ============= Visualization ============= - def _on_viz_thread(self, fn: Any, *args: Any, **kwargs: Any) -> Any: - """Run fn on the Meshcat creator thread. Re-entrant safe.""" - if self._viz_executor is None or current_thread() is self._viz_thread: - return fn(*args, **kwargs) - return self._viz_executor.submit(fn, *args, **kwargs).result() - def get_visualization_url(self) -> str | None: """Get visualization URL if enabled.""" if self._meshcat is not None: - url: str = self._on_viz_thread(self._meshcat.web_url) - return url + return self._meshcat.web_url() return None def publish_visualization(self, ctx: Context | None = None) -> None: - """Publish current state to visualization. - - Args: - ctx: Context to publish. Uses live context if None. - """ - if self._meshcat_visualizer is None: + """Publish current state to visualization.""" + if self._meshcat_visualizer is None or self._meshcat is None: return - if ctx is None: ctx = self._live_context - if ctx is not None: viz_ctx = self._diagram.GetSubsystemContext(self._meshcat_visualizer, ctx) - self._on_viz_thread(self._meshcat_visualizer.ForcedPublish, viz_ctx) + self._meshcat.forced_publish(self._meshcat_visualizer, viz_ctx) def _set_preview_positions( self, plant_ctx: Context, robot_id: WorldRobotID, positions: NDArray[np.float64] @@ -957,7 +982,7 @@ def show_preview(self, robot_id: WorldRobotID) -> None: if robot_data is None or robot_data.preview_model_instance is None: return model_name = self._plant.GetModelInstanceName(robot_data.preview_model_instance) - self._on_viz_thread(self._meshcat.SetProperty, f"visualizer/{model_name}", "visible", True) + self._meshcat.SetProperty(f"visualizer/{model_name}", "visible", True) def hide_preview(self, robot_id: WorldRobotID) -> None: """Hide the preview (yellow ghost) robot in Meshcat.""" @@ -967,7 +992,7 @@ def hide_preview(self, robot_id: WorldRobotID) -> None: if robot_data is None or robot_data.preview_model_instance is None: return model_name = self._plant.GetModelInstanceName(robot_data.preview_model_instance) - self._on_viz_thread(self._meshcat.SetProperty, f"visualizer/{model_name}", "visible", False) + self._meshcat.SetProperty(f"visualizer/{model_name}", "visible", False) def animate_path( self, @@ -999,10 +1024,9 @@ def animate_path( time.sleep(dt) def close(self) -> None: - """Shut down the viz executor thread.""" - if self._viz_executor is not None: - self._viz_executor.shutdown(wait=False) - self._viz_executor = None + """Shut down the viz thread.""" + if self._meshcat is not None: + self._meshcat.close() # ============= Direct Access (use with caution) =============