diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index e24e425460..b2a0be40d4 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -108,18 +108,32 @@ class NativeModuleConfig(ModuleConfig): cli_exclude: frozenset[str] = frozenset() cli_name_override: dict[str, str] = Field(default_factory=dict) + def _native_ignore_fields(self) -> set[str]: + """Inherited NativeModuleConfig fields *not* redeclared in any subclass. + + A subclass that redeclares an inherited field (e.g. ``frame_id``) is + signalling it wants that field exposed — usually as a CLI arg or + stdin config entry — so we don't filter it out. + """ + ignore = set(NativeModuleConfig.model_fields) + for klass in self.__class__.__mro__: + if klass is NativeModuleConfig: + break + ignore -= set(getattr(klass, "__annotations__", {})) + return ignore + def to_config_dict(self) -> dict[str, Any]: """ Return module-specific config fields as a plain dict (for stdin JSON). """ - ignore_fields = set(NativeModuleConfig.model_fields) + ignore_fields = self._native_ignore_fields() return { k: v for k, v in self.model_dump().items() if k not in ignore_fields and v is not None } def to_cli_args(self) -> list[str]: """Convert subclass config fields to CLI args (--name value).""" - ignore_fields = {f for f in NativeModuleConfig.model_fields if f != "frame_id"} + ignore_fields = self._native_ignore_fields() args: list[str] = [] for f in self.__class__.model_fields: if f in ignore_fields: diff --git a/dimos/hardware/sensors/lidar/fastlio2/fastlio_blueprints.py b/dimos/hardware/sensors/lidar/fastlio2/fastlio_blueprints.py index 89f3e82ab8..f09b3e5822 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/fastlio_blueprints.py +++ b/dimos/hardware/sensors/lidar/fastlio2/fastlio_blueprints.py @@ -12,20 +12,108 @@ # See the License for the specific language governing permissions and # limitations under the License. +from pathlib import Path +from typing import TYPE_CHECKING + +from reactivex.disposable import Disposable from dimos.core.coordination.blueprints import autoconnect +from dimos.core.core import rpc +from dimos.core.stream import In, Out from dimos.hardware.sensors.lidar.fastlio2.module import FastLio2 from dimos.mapping.voxels import VoxelGridMapper +from dimos.memory2.module import MemoryModule, MemoryModuleConfig, Recorder, RecorderConfig +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.utils.testing.replay import timed_playback from dimos.visualization.vis_module import vis_module +if TYPE_CHECKING: + from rerun._baseclasses import Archetype + + +class FastlioMemoryConfig(RecorderConfig): + db_path: str | Path = "recording_fastlio.db" + default_frame_id: str = "base_link" + + voxel_size = 0.05 +class FastlioMemory(Recorder): + config: FastlioMemoryConfig + lidar: In[PointCloud2] + odometry: In[Odometry] + + @rpc + def start(self) -> None: + super().start() + + def _on_odom(msg: Odometry) -> None: + self.tf.publish(Transform.from_odometry(msg)) + + self.register_disposable(Disposable(self.odometry.subscribe(_on_odom))) + + +class FastlioReplayConfig(MemoryModuleConfig): + db_path: str | Path = "recording_fastlio.db" + speed: float = 1.0 + + +class FastlioReplay(MemoryModule): + """Replays a FastLIO2 recording (lidar + odometry) at real-time speed. + + Drop-in replacement for ``FastLio2`` when feeding rerun off a recorded session. + Publishes odometry to tf so downstream visualizers see robot pose. + """ + + config: FastlioReplayConfig + lidar: Out[PointCloud2] + odometry: Out[Odometry] + + @rpc + def start(self) -> None: + super().start() + + lidar_stream = self.store.stream("lidar", PointCloud2) + odom_stream = self.store.stream("odometry", Odometry) + + def _publish_odom(msg: Odometry) -> None: + self.tf.publish(Transform.from_odometry(msg)) + self.odometry.publish(msg) + + speed = self.config.speed + + self.register_disposable( + timed_playback( + lambda: ((obs.ts, obs.data) for obs in lidar_stream), + speed=speed, + ).subscribe(self.lidar.publish) + ) + self.register_disposable( + timed_playback( + lambda: ((obs.ts, obs.data) for obs in odom_stream), + speed=speed, + ).subscribe(_publish_odom) + ) + + +def _convert_global_map(msg: PointCloud2) -> "Archetype": + return msg.to_rerun(voxel_size=voxel_size) + + mid360_fastlio = autoconnect( FastLio2.blueprint(voxel_size=voxel_size, map_voxel_size=voxel_size, map_freq=-1), vis_module("rerun"), ).global_config(n_workers=2, robot_model="mid360_fastlio2") +mid360_fastlio_memory = autoconnect( + FastLio2.blueprint(voxel_size=voxel_size, map_voxel_size=voxel_size, map_freq=-1), + vis_module("rerun"), + FastlioMemory.blueprint(), +).global_config(n_workers=3, robot_model="mid360_fastlio2_memory") + mid360_fastlio_voxels = autoconnect( FastLio2.blueprint(), VoxelGridMapper.blueprint(voxel_size=voxel_size, carve_columns=False), @@ -39,6 +127,31 @@ ), ).global_config(n_workers=3, robot_model="mid360_fastlio2_voxels") +mid360_fastlio_replay = autoconnect( + FastlioReplay.blueprint(), + vis_module( + "rerun", + rerun_config={ + "visual_override": { + "world/global_map": _convert_global_map, + }, + }, + ), +).global_config(n_workers=2, robot_model="mid360_fastlio2_replay") + +mid360_fastlio_replay_voxels = autoconnect( + FastlioReplay.blueprint(), + VoxelGridMapper.blueprint(voxel_size=voxel_size, carve_columns=True), + vis_module( + "rerun", + rerun_config={ + "visual_override": { + "world/global_map": _convert_global_map, + }, + }, + ), +).global_config(n_workers=2, robot_model="mid360_fastlio2_replay") + mid360_fastlio_voxels_native = autoconnect( FastLio2.blueprint(voxel_size=voxel_size, map_voxel_size=voxel_size, map_freq=3.0), vis_module( diff --git a/dimos/hardware/sensors/lidar/fastlio2/module.py b/dimos/hardware/sensors/lidar/fastlio2/module.py index b4f624dd56..f7947e2529 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/module.py +++ b/dimos/hardware/sensors/lidar/fastlio2/module.py @@ -125,7 +125,7 @@ class FastLio2Config(NativeModuleConfig): host_imu_data_port: int = SDK_HOST_IMU_DATA_PORT host_log_data_port: int = SDK_HOST_LOG_DATA_PORT - # Resolved in __post_init__, passed as --config_path to the binary + # Resolved from `config` in model_post_init, passed as --config_path to the binary config_path: str | None = None # init_pose is computed from mount; config is resolved to config_path diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index b584553bae..c467b68ead 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -305,7 +305,22 @@ def _port_to_stream(self, name: str, input_topic: In[Any], stream: Stream[Any]) def on_msg(msg: Any) -> None: ts = getattr(msg, "ts", None) or time.time() - frame_id = getattr(msg, "frame_id", None) or default_frame_id + # For msgs that carry a parent→child transform (Odometry, + # TransformStamped), child_frame_id is the body whose pose we + # want to anchor; frame_id is just the parent (often world). + # Plain stamped msgs only have frame_id (the frame the data is in); + # if that's already 'world' the data carries no robot-pose info, + # so fall back to default_frame_id to still anchor it to the + # robot's world pose at this timestamp. + frame_id = ( + getattr(msg, "child_frame_id", None) + or getattr(msg, "frame_id", None) + or default_frame_id + ) + + if frame_id == "world": + frame_id = default_frame_id + transform = self.tf.get("world", frame_id, time_point=ts, time_tolerance=tf_tolerance) pose = transform.to_pose() if transform is not None else None diff --git a/dimos/msgs/geometry_msgs/Transform.py b/dimos/msgs/geometry_msgs/Transform.py index 9b08c8dadd..14d5b0c742 100644 --- a/dimos/msgs/geometry_msgs/Transform.py +++ b/dimos/msgs/geometry_msgs/Transform.py @@ -21,6 +21,7 @@ import rerun as rr from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped + from dimos.msgs.nav_msgs.Odometry import Odometry from dimos_lcm.geometry_msgs import ( Transform as LCMTransform, @@ -161,6 +162,17 @@ def __neg__(self) -> Transform: """Unary minus operator returns the inverse transform.""" return self.inverse() + @classmethod + def from_odometry(cls, odom: Odometry) -> Transform: # type: ignore[name-defined] + """Create a Transform from an Odometry message using its own frame names.""" + return cls( + translation=odom.pose.position, + rotation=odom.pose.orientation, + frame_id=odom.frame_id, + child_frame_id=odom.child_frame_id, + ts=odom.ts, + ) + @classmethod def from_pose(cls, frame_id: str, pose: Pose | PoseStamped) -> Transform: # type: ignore[name-defined] """Create a Transform from a Pose or PoseStamped. diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index 6e25af7704..c9110c3b07 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -155,6 +155,14 @@ def get_transform( return None def get(self, *args, **kwargs) -> Transform | None: # type: ignore[no-untyped-def] + parent_frame = args[0] if args else kwargs.get("parent_frame") + child_frame = args[1] if len(args) > 1 else kwargs.get("child_frame") + if parent_frame is not None and parent_frame == child_frame: + raise ValueError( + f"tf.get() called with same parent and child frame {parent_frame!r}; " + "this is almost always a caller bug — the data is already in that frame" + ) + simple = self.get_transform(*args, **kwargs) if simple is not None: return simple diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index ef7214811c..ac0e571fb6 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -63,6 +63,9 @@ "keyboard-teleop-xarm7": "dimos.robot.manipulators.xarm.blueprints:keyboard_teleop_xarm7", "mid360": "dimos.hardware.sensors.lidar.livox.livox_blueprints:mid360", "mid360-fastlio": "dimos.hardware.sensors.lidar.fastlio2.fastlio_blueprints:mid360_fastlio", + "mid360-fastlio-memory": "dimos.hardware.sensors.lidar.fastlio2.fastlio_blueprints:mid360_fastlio_memory", + "mid360-fastlio-replay": "dimos.hardware.sensors.lidar.fastlio2.fastlio_blueprints:mid360_fastlio_replay", + "mid360-fastlio-replay-voxels": "dimos.hardware.sensors.lidar.fastlio2.fastlio_blueprints:mid360_fastlio_replay_voxels", "mid360-fastlio-voxels": "dimos.hardware.sensors.lidar.fastlio2.fastlio_blueprints:mid360_fastlio_voxels", "mid360-fastlio-voxels-native": "dimos.hardware.sensors.lidar.fastlio2.fastlio_blueprints:mid360_fastlio_voxels_native", "openarm-mock-planner-coordinator": "dimos.robot.manipulators.openarm.blueprints:openarm_mock_planner_coordinator", @@ -136,6 +139,8 @@ "emitter-module": "dimos.utils.demo_image_encoding.EmitterModule", "far-planner": "dimos.navigation.nav_stack.modules.far_planner.far_planner.FarPlanner", "fast-lio2": "dimos.hardware.sensors.lidar.fastlio2.module.FastLio2", + "fastlio-memory": "dimos.hardware.sensors.lidar.fastlio2.fastlio_blueprints.FastlioMemory", + "fastlio-replay": "dimos.hardware.sensors.lidar.fastlio2.fastlio_blueprints.FastlioReplay", "foxglove-bridge": "dimos.robot.foxglove_bridge.FoxgloveBridge", "g1-connection": "dimos.robot.unitree.g1.connection.G1Connection", "g1-connection-base": "dimos.robot.unitree.g1.connection.G1ConnectionBase",