go2 mem2 recorder fixes, misc fixes, new recording

dimos/memory2/module.py

@@ -314,9 +314,8 @@ def on_msg(msg: Any) -> None:
                     "[%s] No tf available for frame '%s' at time %s (msg ts: %s), storing without pose",
                     name,
                     frame_id,
-                    ts,
                     getattr(msg, "ts", None),
                 )
-            stream.append(msg, ts=ts, pose=pose)
+            stream.append(msg, ts=time.time(), pose=pose)
 
         self.register_disposable(Disposable(input_topic.subscribe(on_msg)))

dimos/protocol/tf/tf.py

@@ -154,12 +154,21 @@ def get_transform(
 
         return None
 
-    def get(self, *args, **kwargs) -> Transform | None:  # type: ignore[no-untyped-def]
-        simple = self.get_transform(*args, **kwargs)
+    def get(
+        self,
+        parent_frame: str,
+        child_frame: str,
+        time_point: float | None = None,
+        time_tolerance: float | None = None,
+    ) -> Transform | None:
+        if parent_frame == child_frame:
+            return Transform(frame_id=parent_frame, child_frame_id=child_frame)
+
+        simple = self.get_transform(parent_frame, child_frame, time_point, time_tolerance)
         if simple is not None:
             return simple
 
-        complex = self.get_transform_search(*args, **kwargs)
+        complex = self.get_transform_search(parent_frame, child_frame, time_point, time_tolerance)
 
         if complex is None:
             return None

dimos/robot/cli/topic.py

@@ -37,11 +37,11 @@
 def _resolve_type(type_name: str) -> type:
     for module_name in _modules_to_try:
         try:
-            module = importlib.import_module(module_name)
-            if hasattr(module, type_name):
-                return getattr(module, type_name)  # type: ignore[no-any-return]
+            module = importlib.import_module(f"{module_name}.{type_name}")
         except ImportError:
             continue
+        if hasattr(module, type_name):
+            return getattr(module, type_name)  # type: ignore[no-any-return]
 
     raise ValueError(f"Could not find type '{type_name}' in any known message modules")
 
@@ -81,7 +81,7 @@ def _on_message(msg: object) -> None:
     def on_msg(channel: str, data: bytes) -> None:
         _, msg_name = channel.split("#", 1)  # e.g. "nav_msgs.Odometry"
         pkg, cls_name = msg_name.split(".", 1)  # "nav_msgs", "Odometry"
-        module = importlib.import_module(f"dimos.msgs.{pkg}")
+        module = importlib.import_module(f"dimos.msgs.{pkg}.{cls_name}")
         cls = getattr(module, cls_name)
         print(cls.lcm_decode(data))
 

dimos/robot/unitree/connection.py

@@ -303,6 +303,19 @@ def set_obstacle_avoidance(self, enabled: bool = True) -> None:
             {"api_id": 1001, "parameter": {"enable": int(enabled)}},
         )
 
+    def set_lidar(self, enabled: bool) -> None:
+        """Turn the L1 lidar on or off. The Go2 boots with lidar on; turning it
+        off saves power and stops the spinning sound."""
+
+        def publish() -> None:
+            # rt/utlidar/switch takes a std_msgs::String of "ON" / "OFF"
+            self.conn.datachannel.pub_sub.publish_without_callback(
+                RTC_TOPIC["ULIDAR_SWITCH"],
+                data="ON" if enabled else "OFF",
+            )
+
+        self.loop.call_soon_threadsafe(publish)
+
     def free_walk(self) -> bool:
         """Activate FreeWalk locomotion mode — enables walking and velocity commands."""
         return bool(self.publish_request(RTC_TOPIC["SPORT_MOD"], {"api_id": SPORT_CMD["FreeWalk"]}))

dimos/robot/unitree/dimsim_connection.py

@@ -94,6 +94,9 @@ def balance_stand(self) -> bool:
     def set_obstacle_avoidance(self, enabled: bool = True) -> None:
         pass
 
+    def set_lidar(self, enabled: bool) -> None:
+        pass
+
     def enable_rage_mode(self) -> bool:
         return True
 

dimos/robot/unitree/go2/cli/go2tool.py

@@ -162,12 +162,14 @@ def _on_device(d: Go2Device) -> None:
                 typer.echo("Found:")
                 for i, d in enumerate(devices, 1):
                     typer.echo(f"  {i}. {d.name} ({d.address})")
-                default = "1" if len(devices) == 1 else None
-                idx = typer.prompt("Select device", default=default, type=int)
-                if not 1 <= idx <= len(devices):
-                    typer.echo("Invalid selection.", err=True)
-                    raise typer.Exit(1)
-                target = devices[idx - 1].address
+                if len(devices) == 1:
+                    target = devices[0].address
+                else:
+                    idx = typer.prompt("Select device", type=int)
+                    if not 1 <= idx <= len(devices):
+                        typer.echo("Invalid selection.", err=True)
+                        raise typer.Exit(1)
+                    target = devices[idx - 1].address
 
         wifi_ssid = ssid if ssid is not None else typer.prompt("Wi-Fi SSID")
         wifi_password = (

dimos/robot/unitree/go2/connection.py

@@ -64,6 +64,7 @@ class Go2Mode(str, Enum):
 class ConnectionConfig(ModuleConfig):
     ip: str = Field(default_factory=lambda m: m["g"].robot_ip)
     mode: Go2Mode = Go2Mode.DEFAULT
+    lidar_auto_off: bool = False
 
 
 class Go2ConnectionProtocol(Protocol):
@@ -79,6 +80,7 @@ def standup(self) -> bool: ...
     def liedown(self) -> bool: ...
     def balance_stand(self) -> bool: ...
     def set_obstacle_avoidance(self, enabled: bool = True) -> None: ...
+    def set_lidar(self, enabled: bool) -> None: ...
     def enable_rage_mode(self) -> bool: ...
     def publish_request(self, topic: str, data: dict) -> dict: ...  # type: ignore[type-arg]
 
@@ -166,6 +168,9 @@ def balance_stand(self) -> bool:
     def set_obstacle_avoidance(self, enabled: bool = True) -> None:
         pass
 
+    def set_lidar(self, enabled: bool) -> None:
+        pass
+
     def enable_rage_mode(self) -> bool:
         return True
 
@@ -268,10 +273,15 @@ def onimage(image: Image) -> None:
 
         self.connection.set_obstacle_avoidance(self.config.g.obstacle_avoidance)
 
+        self.connection.set_lidar(True)
+
         # self.record("go2_bigoffice")
 
     @rpc
     def stop(self) -> None:
+        if self.config.lidar_auto_off:
+            self.connection.set_lidar(False)
+
         self.liedown()
 
         if self.connection:
@@ -337,6 +347,11 @@ def balance_stand(self) -> bool:
         """Enter BalanceStand: neutral state for switching locomotion modes"""
         return self.connection.balance_stand()
 
+    @rpc
+    def set_lidar(self, enabled: bool) -> None:
+        """Turn the L1 lidar on or off."""
+        self.connection.set_lidar(enabled)
+
     @rpc
     def enable_rage_mode(self) -> bool:
         """Enable Rage Mode (~2.5 m/s forward velocity envelope).

dimos/robot/unitree/mujoco_connection.py

@@ -237,6 +237,9 @@ def balance_stand(self) -> bool:
     def set_obstacle_avoidance(self, enabled: bool = True) -> None:
         pass
 
+    def set_lidar(self, enabled: bool) -> None:
+        pass
+
     def enable_rage_mode(self) -> bool:
         return True