Fix Unitree Go2 (replay and spatial memory)

dimos/core/global_config.py

@@ -19,8 +19,8 @@
 
 class GlobalConfig(BaseSettings):
     robot_ip: str | None = None
-    use_simulation: bool = False
-    use_replay: bool = False
+    simulation: bool = False
+    replay: bool = False
     n_dask_workers: int = 2
 
     model_config = SettingsConfigDict(
@@ -32,8 +32,8 @@ class GlobalConfig(BaseSettings):
 
     @cached_property
     def unitree_connection_type(self) -> str:
-        if self.use_replay:
-            return "fake"
-        if self.use_simulation:
+        if self.replay:
+            return "replay"
+        if self.simulation:
             return "mujoco"
         return "webrtc"

dimos/navigation/bt_navigator/navigator.py

@@ -36,7 +36,7 @@
 from dimos.utils.logging_config import setup_logger
 from dimos.utils.transform_utils import apply_transform
 
-logger = setup_logger("dimos.navigation.bt_navigator")
+logger = setup_logger(__file__)
 
 
 class NavigatorState(Enum):

dimos/perception/spatial_perception.py

@@ -61,7 +61,7 @@ class SpatialMemory(Module):
     """
 
     # LCM inputs
-    image: In[Image] = None
+    color_image: In[Image] = None
 
     def __init__(
         self,
@@ -173,9 +173,6 @@ def __init__(
         self.frame_count: int = 0
         self.stored_frame_count: int = 0
 
-        # For tracking stream subscription
-        self._subscription = None
-
         # List to store robot locations
         self.robot_locations: list[RobotLocation] = []
 
@@ -199,7 +196,7 @@ def set_video(image_msg: Image) -> None:
             else:
                 logger.warning("Received image message without data attribute")
 
-        unsub = self.image.subscribe(set_video)
+        unsub = self.color_image.subscribe(set_video)
         self._disposables.add(Disposable(unsub))
 
         # Start periodic processing using interval
@@ -208,8 +205,6 @@ def set_video(image_msg: Image) -> None:
 
     @rpc
     def stop(self) -> None:
-        self.stop_continuous_processing()
-
         # Save data before shutdown
         self.save()
 
@@ -224,7 +219,6 @@ def _process_frame(self) -> None:
         if self._latest_video_frame is None or tf is None:
             return
 
-        # print("Processing frame for spatial memory...", tf)
         # Create Pose object with position and orientation
         current_pose = tf.to_pose()
 
@@ -265,8 +259,6 @@ def _process_frame(self) -> None:
             # Get euler angles from quaternion orientation for metadata
             euler = tf.rotation.to_euler()
 
-            print(f"Storing frame {frame_id} at position {current_pose}...")
-
             # Create metadata dictionary with primitive types only
             metadata = {
                 "pos_x": float(current_pose.position.x),
@@ -324,72 +316,6 @@ def query_by_location(
         """
         return self.vector_db.query_by_location(x, y, radius, limit)
 
-    def start_continuous_processing(
-        self, video_stream: Observable, get_pose: callable
-    ) -> disposable.Disposable:
-        """
-        Start continuous processing of video frames from an Observable stream.
-
-        Args:
-            video_stream: Observable of video frames
-            get_pose: Callable that returns position and rotation for each frame
-
-        Returns:
-            Disposable subscription that can be used to stop processing
-        """
-        # Stop any existing subscription
-        self.stop_continuous_processing()
-
-        # Map each video frame to include transform data
-        combined_stream = video_stream.pipe(
-            ops.map(lambda video_frame: {"frame": video_frame, **get_pose()}),
-            # Filter out bad transforms
-            ops.filter(
-                lambda data: data.get("position") is not None and data.get("rotation") is not None
-            ),
-        )
-
-        # Process with spatial memory
-        result_stream = self.process_stream(combined_stream)
-
-        # Subscribe to the result stream
-        self._subscription = result_stream.subscribe(
-            on_next=self._on_frame_processed,
-            on_error=lambda e: logger.error(f"Error in spatial memory stream: {e}"),
-            on_completed=lambda: logger.info("Spatial memory stream completed"),
-        )
-
-        logger.info("Continuous spatial memory processing started")
-        return self._subscription
-
-    def stop_continuous_processing(self) -> None:
-        """
-        Stop continuous processing of video frames.
-        """
-        if self._subscription is not None:
-            try:
-                self._subscription.dispose()
-                self._subscription = None
-                logger.info("Stopped continuous spatial memory processing")
-            except Exception as e:
-                logger.error(f"Error stopping spatial memory processing: {e}")
-
-    def _on_frame_processed(self, result: dict[str, Any]) -> None:
-        """
-        Handle updates from the spatial memory processing stream.
-        """
-        # Log successful frame storage (if stored)
-        position = result.get("position")
-        if position is not None:
-            logger.debug(
-                f"Spatial memory updated with frame at ({position[0]:.2f}, {position[1]:.2f}, {position[2]:.2f})"
-            )
-
-        # Periodically save visual memory to disk (e.g., every 100 frames)
-        if self._visual_memory is not None and self.visual_memory_path is not None:
-            if self.stored_frame_count % 100 == 0:
-                self.save()
-
     @rpc
     def save(self) -> bool:
         """
@@ -650,7 +576,7 @@ def deploy(
     camera: spec.Camera,
 ):
     spatial_memory = dimos.deploy(SpatialMemory, db_path="/tmp/spatial_memory_db")
-    spatial_memory.image.connect(camera.image)
+    spatial_memory.color_image.connect(camera.image)
     spatial_memory.start()
     return spatial_memory
 

dimos/robot/cli/README.md

@@ -43,23 +43,23 @@ This tool also initializes the global config and passes it to the blueprint.
 ```python
 class GlobalConfig(BaseSettings):
     robot_ip: str | None = None
-    use_simulation: bool = False
-    use_replay: bool = False
+    simulation: bool = False
+    replay: bool = False
     n_dask_workers: int = 2
 ```
 
 Configuration values can be set from multiple places in order of precedence (later entries override earlier ones):
 
-- Default value defined on GlobalConfig. (`use_simulation = False`)
-- Value defined in `.env` (`USE_SIMULATION=true`)
-- Value in the environment variable (`USE_SIMULATION=true`)
-- Value coming from the CLI (`--use-simulation` or `--no-use-simulation`)
-- Value defined on the blueprint (`blueprint.global_config(use_simulation=True)`)
+- Default value defined on GlobalConfig. (`simulation = False`)
+- Value defined in `.env` (`SIMULATION=true`)
+- Value in the environment variable (`SIMULATION=true`)
+- Value coming from the CLI (`--simulation` or `--no-simulation`)
+- Value defined on the blueprint (`blueprint.global_config(simulation=True)`)
 
 For environment variables/`.env` values, you have to prefix the name with `DIMOS_`.
 
 For the command line, you call it like this:
 
 ```bash
-dimos-robot --use-simulation run unitree-go2
-```
+dimos-robot --simulation run unitree-go2
+```

dimos/robot/unitree_webrtc/depth_module.py

@@ -79,7 +79,7 @@ def __init__(
         self._stop_processing = threading.Event()
 
         if global_config:
-            if global_config.use_simulation:
+            if global_config.simulation:
                 self.gt_depth_scale = 1.0
 
     @rpc

dimos/robot/unitree_webrtc/type/map.py

@@ -52,7 +52,7 @@ def __init__(
         self.max_height = max_height
 
         if global_config:
-            if global_config.use_simulation:
+            if global_config.simulation:
                 self.min_height = 0.3
 
         super().__init__(**kwargs)

dimos/robot/unitree_webrtc/unitree_go2.py

@@ -158,7 +158,7 @@ def __init__(
         cfg = global_config or GlobalConfig()
         self.ip = ip if ip is not None else cfg.robot_ip
         self.connection_type = connection_type or cfg.unitree_connection_type
-        self.rectify_image = not cfg.use_simulation
+        self.rectify_image = not cfg.simulation
         self.tf = TF()
         self.connection = None
 
@@ -205,7 +205,7 @@ def start(self) -> None:
         self.connection.start()
 
         # Connect sensor streams to outputs
-        unsub = self.connection.lidar_stream().subscribe(self.lidar.publish)
+        unsub = self.connection.lidar_stream().subscribe(self._on_lidar)
         self._disposables.add(unsub)
 
         unsub = self.connection.odom_stream().subscribe(self._publish_tf)
@@ -227,16 +227,22 @@ def stop(self) -> None:
             self.connection.stop()
         super().stop()
 
+    def _on_lidar(self, msg: LidarMessage) -> None:
+        if self.lidar.transport:
+            self.lidar.publish(msg)
+
     def _on_video(self, msg: Image) -> None:
         """Handle incoming video frames and publish synchronized camera data."""
         # Apply rectification if enabled
         if self.rectify_image:
             rectified_msg = rectify_image(msg, self.camera_matrix, self.dist_coeffs)
             self._last_image = rectified_msg
-            self.color_image.publish(rectified_msg)
+            if self.color_image.transport:
+                self.color_image.publish(rectified_msg)
         else:
             self._last_image = msg
-            self.color_image.publish(msg)
+            if self.color_image.transport:
+                self.color_image.publish(msg)
 
         # Publish camera info and pose synchronized with video
         timestamp = msg.ts if msg.ts else time.time()
@@ -248,21 +254,34 @@ def _publish_gps_location(self, msg: LatLon) -> None:
 
     def _publish_tf(self, msg) -> None:
         self._odom = msg
-        self.odom.publish(msg)
+        if self.odom.transport:
+            self.odom.publish(msg)
         self.tf.publish(Transform.from_pose("base_link", msg))
+
+        # Publish camera_link transform
         camera_link = Transform(
             translation=Vector3(0.3, 0.0, 0.0),
             rotation=Quaternion(0.0, 0.0, 0.0, 1.0),
             frame_id="base_link",
             child_frame_id="camera_link",
             ts=time.time(),
         )
-        self.tf.publish(camera_link)
+
+        map_to_world = Transform(
+            translation=Vector3(0.0, 0.0, 0.0),
+            rotation=Quaternion(0.0, 0.0, 0.0, 1.0),
+            frame_id="map",
+            child_frame_id="world",
+            ts=time.time(),
+        )
+
+        self.tf.publish(camera_link, map_to_world)
 
     def _publish_camera_info(self, timestamp: float) -> None:
         header = Header(timestamp, "camera_link")
         self.lcm_camera_info.header = header
-        self.camera_info.publish(self.lcm_camera_info)
+        if self.camera_info.transport:
+            self.camera_info.publish(self.lcm_camera_info)
 
     def _publish_camera_pose(self, timestamp: float) -> None:
         """Publish camera pose from TF lookup."""
@@ -282,7 +301,8 @@ def _publish_camera_pose(self, timestamp: float) -> None:
                     position=transform.translation,
                     orientation=transform.rotation,
                 )
-                self.camera_pose.publish(pose_msg)
+                if self.camera_pose.transport:
+                    self.camera_pose.publish(pose_msg)
             else:
                 logger.debug("Could not find transform from world to camera_link")
 
@@ -535,9 +555,6 @@ def _deploy_perception(self) -> None:
         self.spatial_memory_module.color_image.transport = core.pSHMTransport(
             "/go2/color_image", default_capacity=DEFAULT_CAPACITY_COLOR_IMAGE
         )
-        self.spatial_memory_module.odom.transport = core.LCMTransport(
-            "/go2/camera_pose", PoseStamped
-        )
 
         logger.info("Spatial memory module deployed and connected")
 

dimos/web/command-center-extension/package.json

@@ -22,7 +22,7 @@
     "@foxglove/eslint-plugin": "2.1.0",
     "@foxglove/extension": "2.34.0",
     "@types/d3": "^7.4.3",
-    "@types/leaflet": "^1.9.20",
+    "@types/leaflet": "^1.9.21",
     "@types/react": "18.3.24",
     "@types/react-dom": "18.3.7",
     "create-foxglove-extension": "1.0.6",
@@ -35,6 +35,7 @@
   "dependencies": {
     "@types/pako": "^2.0.4",
     "d3": "^7.9.0",
+    "leaflet": "^1.9.4",
     "pako": "^2.1.0",
     "react-leaflet": "^4.2.1",
     "socket.io-client": "^4.8.1"