fix(unitree-go2): minor issues (#1307)

* fix(unitree-go2): minor issues

* fix(unitree-go2): minor issues

* fix stand down/up

* fix mypy

Author: paul-nechifor

dimos/agents/skills/navigation.py

@@ -63,6 +63,7 @@ def __init__(self) -> None:
 
     @rpc
     def start(self) -> None:
+        super().start()
         self._disposables.add(Disposable(self.color_image.subscribe(self._on_color_image)))
         self._disposables.add(Disposable(self.odom.subscribe(self._on_odom)))
         self._skill_started = True
@@ -92,12 +93,12 @@ def tag_location(self, location_name: str) -> str:
 
         if not self._skill_started:
             raise ValueError(f"{self} has not been started.")
-        tf = self.tf.get("map", "base_link", time_tolerance=2.0)
-        if not tf:
-            return "Could not get the robot's current transform."
 
-        position = tf.translation
-        rotation = tf.rotation.to_euler()
+        if not self._latest_odom:
+            return "No odometry data received yet, cannot tag location."
+
+        position = self._latest_odom.position
+        rotation = self._latest_odom.orientation
 
         location = RobotLocation(
             name=location_name,

dimos/core/__init__.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import multiprocessing as mp
-import signal
 import time
 from typing import TYPE_CHECKING, cast
 
@@ -266,31 +265,6 @@ def start(n: int | None = None, memory_limit: str = "auto") -> DimosCluster:
     )
 
     patched_client = patchdask(client, cluster)
-    patched_client._shutting_down = False  # type: ignore[attr-defined]
-
-    # Signal handler with proper exit handling
-    def signal_handler(sig, frame) -> None:  # type: ignore[no-untyped-def]
-        # If already shutting down, force exit
-        if patched_client._shutting_down:  # type: ignore[attr-defined]
-            import os
-
-            console.print("[red]Force exit!")
-            os._exit(1)
-
-        patched_client._shutting_down = True  # type: ignore[attr-defined]
-        console.print(f"[yellow]Shutting down (signal {sig})...")
-
-        try:
-            patched_client.close_all()  # type: ignore[attr-defined]
-        except Exception:
-            pass
-
-        import sys
-
-        sys.exit(0)
-
-    signal.signal(signal.SIGINT, signal_handler)
-    signal.signal(signal.SIGTERM, signal_handler)
 
     return patched_client
 

dimos/e2e_tests/test_person_follow.py

@@ -64,10 +64,8 @@ def test_person_follow(
 ) -> None:
     start_blueprint("--mujoco-start-pos", "-6.18 0.96", "run", "unitree-go2-agentic")
 
-    lcm_spy.save_topic("/rpc/HumanInput/start/res")
-    lcm_spy.wait_for_saved_topic("/rpc/HumanInput/start/res", timeout=120.0)
-    lcm_spy.save_topic("/agent")
-    lcm_spy.wait_for_saved_topic_content("/agent", b"AIMessage", timeout=120.0)
+    lcm_spy.save_topic("/rpc/Agent/on_system_modules/res")
+    lcm_spy.wait_for_saved_topic("/rpc/Agent/on_system_modules/res", timeout=120.0)
 
     time.sleep(5)
 

dimos/e2e_tests/test_spatial_memory.py

@@ -34,10 +34,8 @@ def test_spatial_memory_navigation(
 ) -> None:
     start_blueprint("run", "unitree-go2-agentic")
 
-    lcm_spy.save_topic("/rpc/HumanInput/start/res")
-    lcm_spy.wait_for_saved_topic("/rpc/HumanInput/start/res", timeout=120.0)
-    lcm_spy.save_topic("/agent")
-    lcm_spy.wait_for_saved_topic_content("/agent", b"AIMessage", timeout=120.0)
+    lcm_spy.save_topic("/rpc/Agent/on_system_modules/res")
+    lcm_spy.wait_for_saved_topic("/rpc/Agent/on_system_modules/res", timeout=120.0)
 
     time.sleep(5)
 

dimos/navigation/frontier_exploration/wavefront_frontier_goal_selector.py

@@ -735,6 +735,16 @@ def stop_exploration(self) -> bool:
         ):
             self.exploration_thread.join(timeout=2.0)
 
+        # Publish current location as goal to stop the robot.
+        if self.latest_odometry is not None:
+            goal = PoseStamped(
+                position=self.latest_odometry.position,
+                orientation=self.latest_odometry.orientation,
+                frame_id="world",
+                ts=self.latest_odometry.ts,
+            )
+            self.goal_request.publish(goal)
+
         logger.info("Stopped autonomous frontier exploration")
         return True
 

dimos/robot/unitree/connection.py

@@ -17,7 +17,7 @@
 import functools
 import threading
 import time
-from typing import TypeAlias
+from typing import Any, TypeAlias
 
 import numpy as np
 from numpy.typing import NDArray
@@ -34,7 +34,6 @@
     WebRTCConnectionMethod,
 )
 
-from dimos.core import rpc
 from dimos.core.resource import Resource
 from dimos.msgs.geometry_msgs import Pose, Transform, Twist
 from dimos.msgs.sensor_msgs import Image, PointCloud2
@@ -228,7 +227,7 @@ def run_unsubscription() -> None:
         )
 
     # Generic sync API call (we jump into the client thread)
-    def publish_request(self, topic: str, data: dict):  # type: ignore[no-untyped-def, type-arg]
+    def publish_request(self, topic: str, data: dict[Any, Any]) -> Any:
         future = asyncio.run_coroutine_threadsafe(
             self.conn.datachannel.pub_sub.publish_request_new(topic, data), self.loop
         )
@@ -275,33 +274,20 @@ def video_stream(self) -> Observable[Image]:
     def lowstate_stream(self) -> Observable[LowStateMsg]:
         return backpressure(self.unitree_sub_stream(RTC_TOPIC["LOW_STATE"]))
 
-    def standup_ai(self) -> bool:
-        return self.publish_request(RTC_TOPIC["SPORT_MOD"], {"api_id": SPORT_CMD["BalanceStand"]})  # type: ignore[no-any-return]
-
-    def standup_normal(self) -> bool:
-        self.publish_request(RTC_TOPIC["SPORT_MOD"], {"api_id": SPORT_CMD["StandUp"]})
-        time.sleep(0.5)
-        self.publish_request(RTC_TOPIC["SPORT_MOD"], {"api_id": SPORT_CMD["RecoveryStand"]})
-        return True
-
-    @rpc
     def standup(self) -> bool:
-        if self.mode == "ai":
-            return self.standup_ai()
-        else:
-            return self.standup_normal()
+        return bool(self.publish_request(RTC_TOPIC["SPORT_MOD"], {"api_id": SPORT_CMD["StandUp"]}))
 
-    @rpc
     def liedown(self) -> bool:
-        return self.publish_request(RTC_TOPIC["SPORT_MOD"], {"api_id": SPORT_CMD["StandDown"]})  # type: ignore[no-any-return]
+        return bool(
+            self.publish_request(RTC_TOPIC["SPORT_MOD"], {"api_id": SPORT_CMD["StandDown"]})
+        )
 
     async def handstand(self):  # type: ignore[no-untyped-def]
         return self.publish_request(
             RTC_TOPIC["SPORT_MOD"],
             {"api_id": SPORT_CMD["Standup"], "parameter": {"data": True}},
         )
 
-    @rpc
     def color(self, color: VUI_COLOR = VUI_COLOR.RED, colortime: int = 60) -> bool:
         return self.publish_request(  # type: ignore[no-any-return]
             RTC_TOPIC["VUI"],

dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py

@@ -52,7 +52,7 @@
             image_topic="/world/color_image",
             optical_frame="camera_optical",
         ),
-        "world/global_map": lambda grid: grid.to_rerun(voxel_size=0.1),
+        "world/global_map": lambda grid: grid.to_rerun(voxel_size=0.1, mode="boxes"),
         "world/navigation_costmap": lambda grid: grid.to_rerun(
             colormap="Accent",
             z_offset=0.015,

dimos/simulation/engines/mujoco_engine.py

@@ -21,7 +21,7 @@
 from typing import TYPE_CHECKING
 
 import mujoco
-import mujoco.viewer as viewer  # type: ignore[import-untyped]
+import mujoco.viewer as viewer  # type: ignore[import-untyped,import-not-found]
 
 from dimos.simulation.engines.base import SimulationEngine
 from dimos.simulation.utils.xml_parser import JointMapping, build_joint_mappings

pyproject.toml

@@ -360,6 +360,7 @@ module = [
     "mujoco_playground.*",
     "nav_msgs.*",
     "open_clip",
+    "pinocchio",
     "piper_sdk.*",
     "plotext",
     "plum.*",
@@ -375,7 +376,6 @@ module = [
     "torchreid",
     "ultralytics.*",
     "unitree_webrtc_connect.*",
-    "watchdog.*",
     "xarm.*",
 ]
 ignore_missing_imports = true