diff --git a/dimos/hardware/sensors/lidar/pointlio/mount_correction.py b/dimos/hardware/sensors/lidar/pointlio/mount_correction.py index 6f99398faf..c962ab9ee7 100644 --- a/dimos/hardware/sensors/lidar/pointlio/mount_correction.py +++ b/dimos/hardware/sensors/lidar/pointlio/mount_correction.py @@ -52,7 +52,7 @@ def _transform_to_matrix(transform: Transform) -> np.ndarray: return matrix -def _base_to_frame_matrix(loader: UrdfLoader, leaf_frame: str) -> np.ndarray: +def base_to_frame_matrix(loader: UrdfLoader, leaf_frame: str) -> np.ndarray: """Compose the fixed-joint chain from the model root down to ``leaf_frame``.""" static_transforms = loader.static_transforms chain: list[Transform] = [] @@ -73,9 +73,9 @@ def mount_correction_matrix( sensor_frame: str = "mid360_link", ) -> list[float]: """Row-major 4x4 (16 floats) mapping the ``original`` mount to the ``new`` one.""" - original = _base_to_frame_matrix( + original = base_to_frame_matrix( UrdfLoader(name="original", model_path=original_static_tf), sensor_frame ) - new = _base_to_frame_matrix(UrdfLoader(name="new", model_path=new_static_tf), sensor_frame) + new = base_to_frame_matrix(UrdfLoader(name="new", model_path=new_static_tf), sensor_frame) correction = np.linalg.inv(new) @ original return [float(value) for value in correction.flatten()] diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 61220a0fae..74bb5e0c03 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -117,10 +117,12 @@ "unitree-go2-detection": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2_detection:unitree_go2_detection", "unitree-go2-fleet": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_fleet:unitree_go2_fleet", "unitree-go2-keyboard-teleop": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_keyboard_teleop:unitree_go2_keyboard_teleop", + "unitree-go2-l1-lidar": "dimos.robot.unitree.go2.blueprints.navigation.unitree_go2_l1_lidar:unitree_go2_l1_lidar", "unitree-go2-markers": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_markers", "unitree-go2-memory": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_memory", "unitree-go2-mid360-record": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_mid360_record:unitree_go2_mid360_record", "unitree-go2-nav-3d": "dimos.robot.unitree.go2.blueprints.navigation.unitree_go2_nav_3d:unitree_go2_nav_3d", + "unitree-go2-nav-3d-rotated": "dimos.robot.unitree.go2.blueprints.navigation.unitree_go2_nav_3d:unitree_go2_nav_3d_rotated", "unitree-go2-relocalization": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_relocalization", "unitree-go2-ros": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2_ros:unitree_go2_ros", "unitree-go2-security": "dimos.robot.unitree.go2.blueprints.agentic.unitree_go2_security:unitree_go2_security", @@ -241,6 +243,7 @@ "rerun-web-socket-server": "dimos.visualization.rerun.websocket_server.RerunWebSocketServer", "security-module": "dimos.experimental.security_demo.security_module.SecurityModule", "semantic-search": "dimos.memory2.module.SemanticSearch", + "simple-lidar": "dimos.robot.unitree.go2.simple_lidar.SimpleLidar", "simple-phone-teleop": "dimos.teleop.phone.phone_extensions.SimplePhoneTeleop", "simple-planner": "dimos.navigation.cmu_nav.modules.simple_planner.simple_planner.SimplePlanner", "spatial-memory": "dimos.perception.spatial_perception.SpatialMemory", diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record_combined.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record_combined.py new file mode 100644 index 0000000000..d8c0f860df --- /dev/null +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record_combined.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +# Copyright 2026 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. + +"""Drive-and-record blueprint for the Go2 + Mid-360 rig. + +Pygame WASD teleop drives the dog while Point-LIO odom+lidar, the Go2's lidar/odom, +and the front camera are recorded into a memory2 db. The Go2/Mid-360 mount frames are +published continuously onto tf so they're captured in the recording. Raw Livox capture +is opt-in: set ``RECORD_PCAP=1`` to also record a .pcap of the Mid-360 UDP stream. + +The lidar IP comes from ``DIMOS_MID360_LIDAR_IP`` (shared by the Mid-360 / pcap +capture and Point-LIO). Run it for a timestamped ``recordings/`` folder:: + + export DIMOS_MID360_LIDAR_IP=192.168.1.171 + uv run python dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record_combined.py +""" + +from datetime import datetime +import os +from pathlib import Path + +import numpy as np + +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.coordination.module_coordinator import ModuleCoordinator +from dimos.core.global_config import global_config +from dimos.hardware.sensors.lidar.livox.module import Mid360 +from dimos.hardware.sensors.lidar.pointlio.module import PointLio +from dimos.hardware.sensors.lidar.pointlio.mount_correction import base_to_frame_matrix +from dimos.hardware.sensors.lidar.virtual_mid360.recorder import Mid360PcapRecorder +from dimos.navigation.cmu_nav.frames import FRAME_ODOM +from dimos.navigation.movement_manager.movement_manager import MovementManager +from dimos.robot.unitree.go2.config import mid360_urdf_path +from dimos.robot.unitree.go2.go2_mid360_recorder import Go2Mid360Recorder +from dimos.robot.unitree.go2.go2_mid360_static_transforms import Go2Mid360StaticTf +from dimos.robot.unitree.go2.simple_lidar import SimpleLidar +from dimos.robot.unitree.keyboard_teleop import KeyboardTeleop +from dimos.robot.urdf_loader import UrdfLoader +from dimos.utils.logging_config import set_run_log_dir, setup_logger + +logger = setup_logger() + +# Opt-in raw-Livox pcap capture (default off). Set RECORD_PCAP=1 to include it. +_RECORD_PCAP = os.getenv("RECORD_PCAP", "").lower() in ("1", "true", "yes", "on") + +_TELEOP_LINEAR_SPEED = 0.3 +_TELEOP_ANGULAR_SPEED = 0.6 + +MID360_FRAME = "mid360_link" + +# SimpleLidar lands the Go2 lidar in base_link; re-express it in the Mid-360 frame +# (inv(base_link -> mid360_link) from the mount URDF) so it lines up with Point-LIO's +# cloud on the shared topic. +_base_to_mid360 = base_to_frame_matrix( + UrdfLoader(name="go2_mid360_record", model_path=mid360_urdf_path), MID360_FRAME +) +_lidar_in_mid360 = [float(value) for value in np.linalg.inv(_base_to_mid360).flatten()] + + +def _default_recording_dir() -> Path: + # Local time, with the machine's actual zone abbreviation (not a hardcoded PST). + now = datetime.now().astimezone() + stamp = ( + now.strftime("%Y-%m-%d") + "_" + now.strftime("%I-%M%p").lower() + "-" + now.strftime("%Z") + ) + return Path("recordings") / stamp + + +unitree_go2_mid360_record_combined = autoconnect( + MovementManager.blueprint(), + # hack: pretend the go2 lidar is an extension of the mid360 lidar (transform it into that frame) + SimpleLidar.blueprint(transform=_lidar_in_mid360, output_frame=MID360_FRAME).remappings( + [ + (SimpleLidar, "odom", "go2_odom"), + (SimpleLidar, "lidar", "pointlio_lidar"), + ] + ), + Mid360.blueprint().remappings( + [ + (Mid360, "lidar", "livox_lidar"), + (Mid360, "imu", "livox_imu"), + ] + ), + PointLio.blueprint(frame_mapping={FRAME_ODOM: "world"}).remappings( + [ + (PointLio, "lidar", "pointlio_lidar"), + (PointLio, "odometry", "pointlio_odometry"), + ] + ), + Go2Mid360Recorder.blueprint(), + # Continuously republishes the rig's mount frames onto tf (no latched static tf). + Go2Mid360StaticTf.blueprint(), + # Pygame keyboard teleop (WASD drive + Q/E strafe). Its cmd_vel feeds + # MovementManager's tele_cmd_vel. + KeyboardTeleop.blueprint( + linear_speed=_TELEOP_LINEAR_SPEED, angular_speed=_TELEOP_ANGULAR_SPEED + ).remappings( + [ + (KeyboardTeleop, "cmd_vel", "tele_cmd_vel"), + ] + ), +).global_config(n_workers=12, robot_model="unitree_go2") + +# Opt-in: also capture a raw .pcap of the Mid-360 UDP stream (RECORD_PCAP=1). +if _RECORD_PCAP: + unitree_go2_mid360_record_combined = autoconnect( + unitree_go2_mid360_record_combined, + Mid360PcapRecorder.blueprint(), + ) + + +if __name__ == "__main__": + recording_dir = _default_recording_dir().resolve() + recording_dir.mkdir(parents=True, exist_ok=True) + set_run_log_dir(recording_dir) + global_config.obstacle_avoidance = False + coordinator = ModuleCoordinator.build( + unitree_go2_mid360_record_combined, + {Go2Mid360Recorder.name: {"db_path": str(recording_dir / "mem2.db")}}, + ) + coordinator.loop() diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_l1_lidar.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_l1_lidar.py new file mode 100644 index 0000000000..2d896b9361 --- /dev/null +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_l1_lidar.py @@ -0,0 +1,50 @@ +#!/usr/bin/env python3 +# Copyright 2026 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. + +"""Go2 onboard (L1) lidar re-expressed in the Mid-360 frame. + +Sibling to ``unitree_go2_nav_3d``. Instead of the physical Mid-360 + Point-LIO, +this runs ``SimpleLidar``: it pulls the Go2's onboard lidar over WebRTC, undoes +the onboard world transform back to ``base_link``, then applies the static +``base_link -> mid360_link`` transform (inverted) so the cloud looks like it was +captured by the Mid-360. Downstream consumers calibrated for the Mid-360 mount +therefore see the L1 lidar in the frame they expect. +""" + +import numpy as np + +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.global_config import global_config +from dimos.hardware.sensors.lidar.pointlio.mount_correction import base_to_frame_matrix +from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_basic import rerun_config +from dimos.robot.unitree.go2.config import mid360_urdf_path +from dimos.robot.unitree.go2.simple_lidar import SimpleLidar +from dimos.robot.urdf_loader import UrdfLoader +from dimos.visualization.vis_module import vis_module + +MID360_FRAME = "mid360_link" + +# The Go2 lidar lands in base_link after SimpleLidar undoes the onboard transform. +# Re-express it in the Mid-360 frame: inv(base_link -> mid360_link) maps base_link +# coordinates into mid360_link, and the cloud is stamped mid360_link to match. +_base_to_mid360 = base_to_frame_matrix( + UrdfLoader(name="go2_mid360", model_path=mid360_urdf_path), MID360_FRAME +) +_lidar_in_mid360 = [float(value) for value in np.linalg.inv(_base_to_mid360).flatten()] + +unitree_go2_l1_lidar = autoconnect( + vis_module(viewer_backend=global_config.viewer, rerun_config=rerun_config), + SimpleLidar.blueprint(transform=_lidar_in_mid360, output_frame=MID360_FRAME), +).global_config(n_workers=4, robot_model="unitree_go2") diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index a816e8590d..120a6ae218 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -17,7 +17,7 @@ from typing import Any -from dimos.core.coordination.blueprints import autoconnect +from dimos.core.coordination.blueprints import Blueprint, autoconnect from dimos.core.global_config import global_config from dimos.hardware.sensors.lidar.pointlio.module import PointLio from dimos.hardware.sensors.lidar.pointlio.mount_correction import mount_correction_matrix @@ -34,21 +34,15 @@ from dimos.robot.urdf_loader import UrdfLoader from dimos.visualization.vis_module import vis_module -# Flip on only when the Mid-360 is physically mounted rotated; this un-rotates its -# Point-LIO cloud + odometry back into the normal mount frame. Off = sensor mounted -# normally, no correction (identity). -USE_ROTATED_MID360_MOUNT = False - -_mount_correction = ( - mount_correction_matrix( - original_static_tf=mid360_rotated_urdf_path, new_static_tf=mid360_urdf_path - ) - if USE_ROTATED_MID360_MOUNT - else None +# Correction that un-rotates a physically-rotated Mid-360's Point-LIO cloud + +# odometry back into the normal mount frame. Used by the `_rotated` blueprint; the +# plain blueprint passes None (identity — sensor mounted normally). +_rotated_mount_correction = mount_correction_matrix( + original_static_tf=mid360_rotated_urdf_path, new_static_tf=mid360_urdf_path ) # GO2Connection is a TfModule; feeding it the normal mount publishes those frames -# on its static interval, matching the cloud the hack emits. +# on its static interval, matching the corrected cloud both blueprints emit. go2_mid360_model = UrdfLoader(name="go2_mid360", model_path=mid360_urdf_path) voxel_size = 0.08 @@ -126,65 +120,81 @@ def _static_robot_body(rr: Any) -> list[Any]: }, } -unitree_go2_nav_3d = autoconnect( - vis_module(viewer_backend=global_config.viewer, rerun_config=_nav_rerun_config), - # "mcf" for stair traversal - GO2Connection.blueprint( - static_transforms=dict(go2_mid360_model.static_transforms), - lidar=False, - camera=False, - motion_mode="mcf", - ).remappings( - [ - (GO2Connection, "lidar", "lidar_l1"), - (GO2Connection, "odom", "odom_go2"), - ] - ), - # gravity_align is off so pointlio runs in the raw mount frame. When the rotated - # mount is enabled, transform rewrites its cloud + odometry (and odom->body TF) - # into the normal mount; otherwise transform is identity. auto_build recompiles - # the native binary for the transform support. - PointLio.blueprint( - gravity_align=False, - space_down_sample=False, - auto_build=True, - transform=_mount_correction, - ), - # Record the corrected (normal-mount) cloud + odometry, not PointLio's raw - # tilted output, so a replay reproduces what the nav stack actually consumed. - # The remap only renames the wire topics to PointLio's lidar/odometry; the - # recorder's ports stay pointlio_lidar/pointlio_odometry, so that's what the - # db streams are named. - PointlioRecorder.blueprint().remappings( - [ - (PointlioRecorder, "pointlio_lidar", "lidar"), - (PointlioRecorder, "pointlio_odometry", "odometry"), - ] - ), - # Raw Livox UDP capture (tcpdump). Pulls the lidar IP + iface from env - # (DIMOS_MID360_LIDAR_IP / DIMOS_MID360_PCAP_IFACE); pcap lands in recordings/. - Mid360PcapRecorder.blueprint(), - RayTracingVoxelMap.blueprint( - voxel_size=voxel_size, - emit_every=1, - global_emit_every=50, - max_health=10, - graze_cos=0.85, - ), - # global_map is remapped off so the planner runs purely on the - # incremental local_map + region_bounds pair. - MLSPlannerNative.blueprint( - world_frame="odom", - voxel_size=voxel_size, - robot_height=go2_lidar_height, - wall_clearance_m=0.2, - wall_buffer_m=0.75, - wall_buffer_weight=100.0, - step_threshold_m=0.16, - step_penalty_weight=1.0, - viz_publish_hz=2.0, - ).remappings([(MLSPlannerNative, "global_map", "global_map_unused")]), - GoalRelay.blueprint(), - BasicPathFollower.blueprint(speed=0.5, heading_gain=0.4, max_angular=0.6), - MovementManager.blueprint(), -).global_config(n_workers=10, robot_model="unitree_go2", obstacle_avoidance=False) + +def _nav_3d(mount_correction: list[float] | None) -> Blueprint: + """Go2 3D nav + recording stack. ``mount_correction`` un-rotates a physically + rotated Mid-360 back into the normal mount frame; None = normal mount (no + correction).""" + return autoconnect( + vis_module(viewer_backend=global_config.viewer, rerun_config=_nav_rerun_config), + # "mcf" for stair traversal + GO2Connection.blueprint( + static_transforms=dict(go2_mid360_model.static_transforms), + lidar=False, + camera=False, + motion_mode="mcf", + ).remappings( + [ + (GO2Connection, "lidar", "lidar_l1"), + (GO2Connection, "odom", "odom_go2"), + ] + ), + # gravity_align is off so pointlio runs in the raw mount frame. For the + # rotated blueprint, transform rewrites its cloud + odometry (and odom->body + # TF) into the normal mount; otherwise transform is None (identity). + # auto_build recompiles the native binary for the transform support. + PointLio.blueprint( + gravity_align=False, + space_down_sample=False, + auto_build=True, + transform=mount_correction, + ), + # Record the corrected (normal-mount) cloud + odometry, not PointLio's raw + # tilted output, so a replay reproduces what the nav stack actually consumed. + # The remap only renames the wire topics to PointLio's lidar/odometry; the + # recorder's ports stay pointlio_lidar/pointlio_odometry, so that's what the + # db streams are named. + PointlioRecorder.blueprint().remappings( + [ + (PointlioRecorder, "pointlio_lidar", "lidar"), + (PointlioRecorder, "pointlio_odometry", "odometry"), + ] + ), + # Raw Livox UDP capture (tcpdump). Pulls the lidar IP + iface from env + # (DIMOS_MID360_LIDAR_IP / DIMOS_MID360_PCAP_IFACE); pcap lands in recordings/. + Mid360PcapRecorder.blueprint(), + RayTracingVoxelMap.blueprint( + voxel_size=voxel_size, + emit_every=1, + global_emit_every=50, + max_health=10, + graze_cos=0.85, + ), + # global_map is remapped off so the planner runs purely on the + # incremental local_map + region_bounds pair. + MLSPlannerNative.blueprint( + world_frame="odom", + voxel_size=voxel_size, + robot_height=go2_lidar_height, + wall_clearance_m=0.2, + wall_buffer_m=0.75, + wall_buffer_weight=100.0, + step_threshold_m=0.16, + step_penalty_weight=1.0, + viz_publish_hz=2.0, + ).remappings([(MLSPlannerNative, "global_map", "global_map_unused")]), + GoalRelay.blueprint(), + BasicPathFollower.blueprint(speed=0.5, heading_gain=0.4, max_angular=0.6), + MovementManager.blueprint(), + ) + + +# Standard mount (level, forward-facing): no correction. +unitree_go2_nav_3d = _nav_3d(None).global_config( + n_workers=10, robot_model="unitree_go2", obstacle_avoidance=False +) + +# Rotated mount: un-rotate the cloud + odometry back into the normal mount frame. +unitree_go2_nav_3d_rotated = _nav_3d(_rotated_mount_correction).global_config( + n_workers=10, robot_model="unitree_go2", obstacle_avoidance=False +) diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index 20a25187a2..76246f5d32 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -244,7 +244,7 @@ def on_image(image: Image) -> None: self._latest_video_frame = image if self.config.lidar: - self.register_disposable(self.connection.lidar_stream().subscribe(self.lidar.publish)) + self.register_disposable(self.connection.lidar_stream().subscribe(self._publish_lidar)) self.register_disposable(self.connection.odom_stream().subscribe(self._publish_tf)) self.register_disposable(self.connection.lowstate_stream().subscribe(self._on_lowstate)) self.register_disposable(Disposable(self.cmd_vel.subscribe(self.move))) @@ -277,6 +277,9 @@ def _on_static_publish(self) -> None: self._camera_info_static.frame_id = self.frame_mapping["camera_optical"] self.camera_info.publish(self._camera_info_static) + def _publish_lidar(self, cloud: PointCloud2) -> None: + self.lidar.publish(cloud) + def _publish_tf(self, msg: PoseStamped) -> None: self.tf.publish( Transform( diff --git a/dimos/robot/unitree/go2/simple_lidar.py b/dimos/robot/unitree/go2/simple_lidar.py new file mode 100644 index 0000000000..7a58084645 --- /dev/null +++ b/dimos/robot/unitree/go2/simple_lidar.py @@ -0,0 +1,114 @@ +# Copyright 2026 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. + +"""GO2Connection variant that re-expresses the Go2's onboard lidar in base_link. + +Drop-in replacement for ``GO2Connection`` — it inherits the whole connection +(movement, odom, camera, tf, skills) and changes only how the lidar cloud is +published. The Go2's onboard stack transforms its lidar into the odom/world +frame (``pointcloud2_from_webrtc_lidar`` stamps it ``world``) and accumulates +scans. This undoes both: it applies the inverse of the robot's current world +pose so points land back in ``base_link``, and with ``un_accumulate`` on it +subtracts the prior cloud (in the stable world frame, where accumulated points +keep identical coordinates) so only the new points are published. + +An optional rigid ``transform`` (row-major 4x4, 16 floats) is applied to the +base_link cloud before publishing — mirrors PointLio's ``transform`` and lets +the cloud be re-expressed as if it came from another sensor (e.g. mid360_link). +""" + +from __future__ import annotations + +import numpy as np +from pydantic import field_validator + +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.robot.unitree.go2.connection import ConnectionConfig, GO2Connection + + +def _rows_not_in(points: np.ndarray, reference: np.ndarray) -> np.ndarray: + """Rows of ``points`` (Nx3) that aren't byte-identical rows of ``reference``.""" + if reference.size == 0: + return points + row_dtype = np.dtype([("", points.dtype)] * points.shape[1]) + points_view = np.ascontiguousarray(points).view(row_dtype).ravel() + reference_view = np.ascontiguousarray(reference).view(row_dtype).ravel() + return points[~np.isin(points_view, reference_view)] + + +class SimpleLidarConfig(ConnectionConfig): + # Optional rigid transform applied to the base_link cloud before publishing: + # row-major 4x4 (16 floats), None = identity. Mirrors PointLio's transform. + transform: list[float] | None = None + # Frame the transformed cloud is stamped with. None keeps the body frame; set + # it to re-express the cloud as if it came from another sensor (e.g. mid360_link). + output_frame: str | None = None + # Subtract the previous (accumulated) cloud and publish only the new points. + un_accumulate: bool = True + + @field_validator("transform") + @classmethod + def _validate_transform(cls, value: list[float] | None) -> list[float] | None: + if value is not None and len(value) != 16: + raise ValueError(f"transform must be a row-major 4x4 (16 floats), got {len(value)}") + return value + + +class SimpleLidar(GO2Connection): + """GO2Connection that publishes the onboard lidar back in the base_link frame.""" + + config: SimpleLidarConfig + + _latest_pose: PoseStamped | None = None + _transform: Transform | None = None + _previous_points: np.ndarray | None = None + + def __init__(self, **kwargs: object) -> None: + super().__init__(**kwargs) + if self.config.transform is not None: + matrix = np.array(self.config.transform, dtype=float).reshape(4, 4) + output_frame = self.config.output_frame or self.frame_mapping["body"] + self._transform = Transform.from_matrix( + matrix, frame_id=output_frame, child_frame_id=output_frame + ) + + def _publish_tf(self, msg: PoseStamped) -> None: + self._latest_pose = msg + super()._publish_tf(msg) + + def _publish_lidar(self, cloud: PointCloud2) -> None: + pose = self._latest_pose + if pose is None: + return + if self.config.un_accumulate: + cloud = self._only_new_points(cloud) + base_frame = self.frame_mapping["body"] + # from_pose gives base_link's pose in world; its inverse maps the world + # cloud back into base_link. + world_to_base = Transform.from_pose(base_frame, pose).inverse() + base_cloud = cloud.transform(world_to_base) + if self._transform is not None: + base_cloud = base_cloud.transform(self._transform) + self.lidar.publish(base_cloud) + + def _only_new_points(self, cloud: PointCloud2) -> PointCloud2: + points = cloud.points_f32() + previous = self._previous_points + self._previous_points = points + if previous is None: + return cloud + new_points = _rows_not_in(points, previous) + return PointCloud2.from_numpy(new_points, frame_id=cloud.frame_id, timestamp=cloud.ts)