custom o3d pointcloud picklers

Author: leshy

dimos/multiprocess/actors2/base.py

@@ -14,13 +14,17 @@
 
 from __future__ import annotations
 
+import copyreg
 import enum
 import inspect
 from typing import Any, Callable, Generic, TypeVar, get_args, get_origin, get_type_hints
 
 from distributed.actor import Actor
 
 import dimos.multiprocess.actors2.colors as colors
+from dimos.multiprocess.actors2.o3dpickle import register_picklers
+
+register_picklers()
 
 T = TypeVar("T")
 
@@ -149,13 +153,12 @@ def __reduce__(self) -> Any:
         return (RemoteOut, (self.type, self.name, self.owner.ref if self.owner else None))
 
     def publish(self, value: T):
-        print("PUB REQ", self.owner, value)
         for sub in self.subscribers:
-            print("PUBLISHING", value, "to", sub)
             sub.owner.receive_msg(sub.name, value)
 
     def subscribe(self, remote_input):
         print(self, "adding remote input to subscribers", remote_input)
+        remote_input.owner._try_bind_worker_client()
         self.subscribers.append(remote_input)
         print(self.subscribers)
 

dimos/multiprocess/actors2/o3dpickle.py

@@ -0,0 +1,38 @@
+# Copyright 2025 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.
+
+import copyreg
+
+import numpy as np
+import open3d as o3d
+
+
+def reduce_external(obj):
+    # Convert Vector3dVector to numpy array for pickling
+    points_array = np.asarray(obj.points)
+    return (reconstruct_pointcloud, (points_array,))
+
+
+def reconstruct_pointcloud(points_array):
+    # Create new PointCloud and assign the points
+    pc = o3d.geometry.PointCloud()
+    pc.points = o3d.utility.Vector3dVector(points_array)
+    return pc
+
+
+def register_picklers():
+    # Register for the actual PointCloud class that gets instantiated
+    # We need to create a dummy PointCloud to get its actual class
+    _dummy_pc = o3d.geometry.PointCloud()
+    copyreg.pickle(_dummy_pc.__class__, reduce_external)

dimos/multiprocess/actors2/test_base.py

@@ -17,6 +17,7 @@
 
 from dimos.multiprocess.actors2.base import In, Module, Out, RemoteOut, module, rpc
 from dimos.multiprocess.actors2.base_dask import dimos
+from dimos.robot.unitree_webrtc.type.lidar import LidarMessage
 from dimos.robot.unitree_webrtc.type.map import Map
 from dimos.robot.unitree_webrtc.type.odometry import Odometry
 from dimos.types.path import Path
@@ -27,33 +28,42 @@
 @module
 class RobotClient(Module):
     odometry: Out[Odometry]
+    lidar: Out[LidarMessage]
 
     def __init__(self):
         self.odometry = Out(Odometry, "odometry", self)
+
         self._stop_event = Event()
         self._thread = None
 
-    @rpc
     def start(self):
         self._thread = Thread(target=self.odomloop)
         self._thread.start()
 
     def odomloop(self):
         odomdata = SensorReplay("raw_odometry_rotate_walk", autocast=Odometry.from_msg)
+        lidardata = SensorReplay("office_lidar", autocast=LidarMessage.from_msg)
+
+        lidariter = lidardata.iterate()
         self._stop_event.clear()
         while not self._stop_event.is_set():
             for odom in odomdata.iterate():
                 if self._stop_event.is_set():
+                    print("Stopping odometry stream")
                     return
+                # print(odom)
                 odom.pubtime = time.perf_counter()
                 self.odometry.publish(odom)
+
+                lidarmsg = next(lidariter)
+                lidarmsg.pubtime = time.perf_counter()
+                self.lidar.publish(lidarmsg)
                 time.sleep(0.1)
 
-    @rpc
     def stop(self):
         self._stop_event.set()
         if self._thread and self._thread.is_alive():
-            self._thread.join(timeout=1.0)
+            self._thread.join(timeout=1.0)  # Wait up to 1 second for clean shutdown
 
 
 @module
@@ -66,18 +76,23 @@ def navigate_to(self, target: Vector) -> bool: ...
     def __init__(
         self,
         target_position: In[Vector],
-        map_stream: In[Map],
+        lidar: In[LidarMessage],
         odometry: In[Odometry],
     ):
         self.target_position = target_position
-        self.map_stream = map_stream
+        self.lidar = lidar
         self.odometry = odometry
         self.target_path = Out(Path, "target_path")
 
     @rpc
     def start(self):
         print("navigation odom stream is, subscribing", self.odometry)
-        self.odometry.subscribe(print)
+        self.odometry.subscribe(
+            lambda msg: print("RCV:", (time.perf_counter() - msg.pubtime) * 1000, msg)
+        )
+        self.lidar.subscribe(
+            lambda msg: print("RCV:", (time.perf_counter() - msg.pubtime) * 1000, msg)
+        )
 
 
 def test_introspection():
@@ -97,46 +112,31 @@ def test_instance_introspection():
     robot = RobotClient()
     print(robot)
 
-    map_stream = Out[Map](Map, "map")
     target_stream = Out[Vector](Vector, "map")
     print("\n")
-    print("map stream", map_stream)
+    print("lidar stream", robot.lidar)
     print("target stream", target_stream)
     print("odom stream", robot.odometry)
 
-    nav = Navigation(target_position=target_stream, map_stream=map_stream, odometry=robot.odometry)
+    nav = Navigation(target_position=target_stream, lidar=robot.lidar, odometry=robot.odometry)
     """Test introspection of the Navigation module."""
     assert hasattr(nav, "inputs")
     assert hasattr(nav, "rpcs")
     print("\n\n\n" + nav.io(), "\n\n")
 
 
-#    nav.start()
-
-
 def test_deployment(dimos):
     robot = dimos.deploy(RobotClient)
-
-    map_stream = RemoteOut[Map](Map, "map")
     target_stream = RemoteOut[Vector](Vector, "map")
-    odom_stream = robot.odometry
+
     print("\n")
-    print("map stream", map_stream)
+    print("lidar stream", robot.lidar)
     print("target stream", target_stream)
-    print("odom stream", odom_stream)
-
-    # # print(type(odom_stream.owner))
-    # # print(type(robot))
-
-    # # out = Out(Odometry, "odometry", robot)
+    print("odom stream", robot.odometry)
 
     nav = dimos.deploy(
-        Navigation, target_position=target_stream, map_stream=map_stream, odometry=odom_stream
+        Navigation, target_position=target_stream, lidar=robot.lidar, odometry=robot.odometry
     )
-
-    # # """Test introspection of the Navigation module."""
-    # # assert hasattr(nav, "inputs")
-    # # assert hasattr(nav, "rpcs")
     print("\n\n\n" + robot.io().result(), "\n")
     print(nav.io().result(), "\n\n")
 

dimos/multiprocess/actors2/test_meta.py

@@ -20,6 +20,7 @@
 
 from dimos.multiprocess.actors2.base import dimos
 from dimos.multiprocess.actors2.meta import ActorReference, In, Out, module, rpc
+from dimos.robot.unitree_webrtc.type.lidar import Lidar
 from dimos.robot.unitree_webrtc.type.map import Map
 from dimos.robot.unitree_webrtc.type.odometry import Odometry
 from dimos.types.path import Path
@@ -55,6 +56,7 @@ def receive_message(self, in_name, message):
 @module
 class RobotClient(Module):
     odometry: Out[Odometry]
+    lidar: Out[Lidar]
 
     def __init__(self):
         self.odometry = Out(Odometry, "odometry", self)
@@ -68,6 +70,9 @@ def start(self):
 
     def odomloop(self):
         odomdata = SensorReplay("raw_odometry_rotate_walk", autocast=Odometry.from_msg)
+        lidardata = SensorReplay("office_lidar", autocast=Lidar.from_msg)
+
+        lidariter = lidardata.iterate()
         self._stop_event.clear()
         while not self._stop_event.is_set():
             for odom in odomdata.iterate():
@@ -77,6 +82,7 @@ def odomloop(self):
                 # print(odom)
                 odom.pubtime = time.perf_counter()
                 self.odometry.publish(odom)
+                self.lidar.publish(next(lidariter))
                 time.sleep(0.1)
 
     def stop(self):

dimos/multiprocess/actors2/test_o3dpickle.py

@@ -0,0 +1,36 @@
+# Copyright 2025 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.
+
+import pickle
+
+from dimos.multiprocess.actors2.o3dpickle import register_picklers
+from dimos.robot.unitree_webrtc.type.lidar import LidarMessage
+from dimos.utils.testing import SensorReplay
+
+register_picklers()
+
+
+def test_enode_decode():
+    lidardata = SensorReplay("office_lidar", autocast=LidarMessage.from_msg)
+    lidarmsg = next(lidardata.iterate())
+
+    binarypc = pickle.dumps(lidarmsg.pointcloud)
+
+    # Test pickling and unpickling
+    binary = pickle.dumps(lidarmsg)
+    lidarmsg2 = pickle.loads(binary)
+
+    # Verify the decoded message has the same properties
+    assert isinstance(lidarmsg2, LidarMessage)
+    assert len(lidarmsg2.pointcloud.points) == len(lidarmsg.pointcloud.points)